Update test

neuml · Apr 19, 2024 · 69b4a9e · 69b4a9e
1 parent bc889bd
commit 69b4a9e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
     DESCRIPTION = "".join([line for line in f if "gh-dark-mode-only" not in line])
 
 # Required dependencies
-install = ["faiss-cpu>=1.7.1.post2", "torch>=1.12.1", "transformers>=4.27.0"]
+install = ["faiss-cpu>=1.7.1.post2", "torch>=1.12.1", "transformers>=4.28.0"]
 
 # Required dependencies that are also transformers dependencies
 install += ["huggingface-hub>=0.9.0", "numpy>=1.18.4", "pyyaml>=5.3", "regex>=2022.8.17"]
@@ -57,7 +57,7 @@
 extras["pipeline-text"] = ["fasttext>=0.9.2", "sentencepiece>=0.1.91"]
 
 extras["pipeline-train"] = [
-    "accelerate>=0.19.0",
+    "accelerate>=0.26.0",
     "bitsandbytes>=0.42.0",
     "onnx>=1.11.0",
     "onnxmltools>=1.9.1",

diff --git a/test/python/testpipeline/testtrainer.py b/test/python/testpipeline/testtrainer.py
@@ -283,7 +283,7 @@ def testRTD(self):
         output = os.path.join(tempfile.gettempdir(), "trainer.rtd")
 
         trainer = HFTrainer()
-        model, _ = trainer("hf-internal-testing/tiny-random-electra", self.data, task="token-detection", output_dir=output)
+        model, _ = trainer("hf-internal-testing/tiny-random-electra", self.data, task="token-detection", save_safetensors=False, output_dir=output)
 
         # Test model completed successfully
         self.assertIsNotNone(model)