Update ul2_tasks.py

ul2_tasks.py

@@ -69,34 +69,46 @@ dataset_shapes = {"train": dataset["train"].num_rows,
 TaskRegistry.add(
     "pretrain_biological_ul2",
     source=seqio.FunctionDataSource(
-        dataset_fn=functools.partial(dataset_fn, dataset=dataset),
+        dataset_fn=functools.partial(
+            dataset_fn, path="Siddharth63/biological_dataset",
+        ),
         splits=("train", "validation"),
         caching_permitted=False,
-        num_input_examples=dataset_shapes,
     ),
     preprocessors=[
         functools.partial(
-            target_to_key, key_map={
-                "inputs": None,
-                "targets": None,
-            }, target_key="targets"),
+            target_to_key,
+            key_map={
+                "inputs": "text",
+                "targets": "text",
+            },
+            target_key="targets",
+        ),
         seqio.preprocessors.tokenize,
         functools.partial(
             ul2_objective,
             shard_ds=False,
             use_prefix_lm_task=True,  # use S-denoising
-            rates=[0.4 / len(R_DENOISER_SPAN_LENGTHS)]*len(R_DENOISER_SPAN_LENGTHS) + [
-                0.4 / len(X_DENOISER_SPAN_LENGTHS)]*len(X_DENOISER_SPAN_LENGTHS) + [0.2],  # equal total 40% rate for both R- and X-denoisers + 20% for S-denoising (suggested at the paper chapter 4.5)
+            rates=[0.4 / len(R_DENOISER_SPAN_LENGTHS)] * len(R_DENOISER_SPAN_LENGTHS)
+            + [0.4 / len(X_DENOISER_SPAN_LENGTHS)] * len(X_DENOISER_SPAN_LENGTHS)
+            + [
+                0.2
+            ],  # equal total 40% rate for both R- and X-denoisers + 20% for S-denoising (suggested at the paper chapter 4.5)
             mean_noise_span_lengths=R_DENOISER_SPAN_LENGTHS + X_DENOISER_SPAN_LENGTHS,
             noise_densities=R_DENOISER_CORRUPT_RATES + X_DENOISER_CORRUPT_RATES,
-            optional_task_prefixes=[R_DENOISER_TOKEN_PREFIX]*len(R_DENOISER_SPAN_LENGTHS) + [
-                X_DENOISER_TOKEN_PREFIX]*len(X_DENOISER_SPAN_LENGTHS) + [S_DENOISER_TOKEN_PREFIX],
+            optional_task_prefixes=[R_DENOISER_TOKEN_PREFIX]
+            * len(R_DENOISER_SPAN_LENGTHS)
+            + [X_DENOISER_TOKEN_PREFIX] * len(X_DENOISER_SPAN_LENGTHS)
+            + [S_DENOISER_TOKEN_PREFIX],
             reserved_for_packing=1,  # make room for task prefix token
         ),
         seqio.preprocessors.append_eos_after_trim,
     ],
-    output_features={"targets": DEFAULT_OUTPUT_FEATURES["targets"]},
-    metric_fns=[metrics.accuracy]
+    output_features={
+        "targets": DEFAULT_OUTPUT_FEATURES["targets"],
+        "inputs": seqio.Feature(vocabulary=vocabulary, add_eos=True),
+    },
+    metric_fns=[metrics.accuracy],
 )