Create a data mixture search for Magi using Pyglove.

PiperOrigin-RevId: 509963072
google · Feb 23, 2023 · 6daf35d · 6daf35d
1 parent 62b13ca
commit 6daf35d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/seqio/experimental.py b/seqio/experimental.py
@@ -141,8 +141,8 @@ def add_fully_cached_mixture(
   return MixtureRegistry.add(
       new_name,
       [
-          (new_t.name, mixture._task_to_rate[old_t.name])  # pylint:disable=protected-access
-          for old_t, new_t in zip(mixture.tasks, new_tasks)
+          (new_t.name, old_ratio)  # pylint:disable=protected-access
+          for old_ratio, new_t in zip(mixture.task_ratios, new_tasks)
       ],
   )
 
@@ -504,7 +504,7 @@ def _sentinel_id(vocabulary, sentinel_num=0):
 
     Args:
       vocabulary: a t5.data.vocabularies.Vocabulary
-      sentinel_num: an optional interger, what sentinel should be returned. By
+      sentinel_num: an optional integer, what sentinel should be returned. By
         default it returns the first sentinel.
 
     Returns: