use util in tokenize_dataset

delphi-suite · May 22, 2024 · 3e2af53 · 3e2af53
1 parent 4e4b969
commit 3e2af53
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scripts/tokenize_dataset.py b/scripts/tokenize_dataset.py
@@ -8,6 +8,7 @@
 from huggingface_hub import HfApi
 from transformers import AutoTokenizer
 
+from delphi import utils
 from delphi.dataset.tokenization import get_tokenized_chunks
 
 if __name__ == "__main__":
@@ -107,7 +108,7 @@
     )
 
     print(f"Tokenizing split='{args.split}'...")
-    split_name = args.split.split("[")[0]
+    split_name = utils.hf_split_to_split_name(args.split)
     for chunk_idx, ds_chunk in enumerate(ds_chunks_it):
         chunk_name = f"{split_name}-{chunk_idx:05}.parquet"
         if args.out_dir: