From 4784c0affd9f59b0a34ff0800c64a91b0d20f881 Mon Sep 17 00:00:00 2001 From: shutcode Date: Tue, 30 Jul 2024 15:00:34 +0800 Subject: [PATCH] Update download_calib_dataset.py --- tools/download_calib_dataset.py | 2 +- tools/download_eval_dataset.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/download_calib_dataset.py b/tools/download_calib_dataset.py index 2fb690e1..3bce0498 100644 --- 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 index 29fa5a77..6bb7424e 100644 --- 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.")