Update ul2_tasks.py

ul2_tasks.py

@@ -62,9 +62,6 @@ def target_to_key(x, key_map, target_key):
     return {**key_map, target_key: x}
 
 
-dataset_shapes = {"train": dataset["train"].num_rows,
-                  "validation": dataset["validation"].num_rows}
-
 TaskRegistry.add(
     "pretrain_medical_ul2",
     source=seqio.FunctionDataSource(
@@ -109,47 +106,3 @@ TaskRegistry.add(
     },
     metric_fns=[metrics.accuracy],
 )
-
-# dataset_name = "gs://medical-siddharth/medical_data"
-# dataset_params = {"from_disk_path": dataset_name}
-
-# if "from_disk_path" in dataset_params:
-#     dataset = load_from_disk(dataset_params.get("from_disk_path"))
-# else:
-#     dataset = load_dataset(**dataset_params)
-
-# dataset_shapes = {"train": dataset["train"].num_rows,
-#                   "validation": dataset["validation"].num_rows}
-
-# TaskRegistry.add(
-#     "pretrain_medical_ul2",
-#     source=seqio.FunctionDataSource(
-#         dataset_fn=functools.partial(dataset_fn, dataset=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"),
-#         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)
-#             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],
-#             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]
-# )