Merge pull request #17 from shutcode/main

Fix download ptb warning
ModelTC · Jul 30, 2024 · ddc9048 · ddc9048
2 parents babad86 + 4784c0a
commit ddc9048
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tools/download_calib_dataset.py b/tools/download_calib_dataset.py
@@ -29,7 +29,7 @@ def download(calib_dataset_name, path):
         calib_dataset.save_to_disk(save_path)
         logger.info("download wikitext2 for calib finished.")
     if "ptb" in calib_dataset_name:
-        calib_dataset = load_dataset("ptb_text_only", "penn_treebank", split="train")
+        calib_dataset = load_dataset("ptb_text_only", "penn_treebank", split="train", trust_remote_code=True)
         save_path = os.path.join(path, "ptb")
         calib_dataset.save_to_disk(save_path)
         logger.info("download ptb for calib finished.")

diff --git a/tools/download_eval_dataset.py b/tools/download_eval_dataset.py
@@ -24,7 +24,7 @@ def download(calib_dataset_name, path):
         calib_dataset.save_to_disk(save_path)
         logger.info("download wikitext2 for eval finished.")
     if "ptb" in calib_dataset_name:
-        calib_dataset = load_dataset("ptb_text_only", "penn_treebank", split="test")
+        calib_dataset = load_dataset("ptb_text_only", "penn_treebank", split="test", trust_remote_code=True)
         save_path = os.path.join(path, "ptb")
         calib_dataset.save_to_disk(save_path)
         logger.info("download ptb for eval finished.")