reduce ut time consumption (#499)

Signed-off-by: Xin He <xin3.he@intel.com>
intel · Oct 23, 2023 · c4f9bfa · c4f9bfa
1 parent 5eda1a3
commit c4f9bfa
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 4 deletions.
diff --git a/tests/requirements.txt b/tests/requirements.txt
@@ -1,6 +1,6 @@
 datasets>=1.17
 mlflow
-nlpaug
+nlpaug==1.1.9
 sacremoses
 neural-compressor
 onnx>=1.10

diff --git a/tests/test_data_augmentation.py b/tests/test_data_augmentation.py
@@ -41,7 +41,7 @@ def test_text_generation_augmentation(self):
         aug = DataAugmentation(augmenter_type="TextGenerationAug")
         aug.input_dataset = self.origin_data
         aug.output_path = os.path.join(self.result_path, "test1.cvs")
-        aug.augmenter_arguments = {'model_name_or_path': 'gpt2-medium'}
+        aug.augmenter_arguments = {'model_name_or_path': 'hf-internal-testing/tiny-random-gpt2'}
         aug.data_augment()
         print("Augmented data:")
         count = 0
@@ -84,7 +84,7 @@ def test_contextualwordembsforsentence_augmentation(self):
         aug.input_dataset = self.origin_data
         aug.column_names = "sentence"
         aug.output_path = os.path.join(self.result_path, "test2.cvs")
-        aug.augmenter_arguments = {"model_path": "xlnet-base-cased"}
+        aug.augmenter_arguments = {"model_path": "hf-internal-testing/tiny-random-xlnet"}
         aug.data_augment()
         raw_datasets = load_dataset("csv", data_files=aug.output_path, delimiter="\t", split="train")
         self.assertTrue(len(raw_datasets) == 10)

diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py
@@ -69,7 +69,6 @@ def test_evaluate_for_JitModel(self):
         )
         self.assertEqual(results["results"]["piqa"]["acc"], 0.6)
 
-
     def test_cnn_daily(self):
         from intel_extension_for_transformers.llm.evaluation.hf_eval import summarization_evaluate
         model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")