update test

nod-ai · Nov 27, 2023 · 05aab43 · 05aab43
1 parent db231ab
commit 05aab43
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/python/turbine_models/tests/llama_test.py b/python/turbine_models/tests/llama_test.py
@@ -13,15 +13,15 @@
 class LLamaTest(unittest.TestCase):
     def testExportTransformerModel(self):
         llama.export_transformer_model(
-            "meta-llama/Llama-2-7b-chat-hf",
+            hf_model_name="meta-llama/Llama-2-7b-chat-hf",
             # TODO: replace with github secret
-            "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",
-            "torch",
-            10,
-            "safetensors",
-            "llama_f32.safetensors",
-            None,
-            "f32",
+            hf_auth_token="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",
+            compile_to="torch",
+            benchmark_steps=10,
+            external_weights="safetensors",
+            external_weight_file="llama_f32.safetensors",
+            quantization=None,
+            precision="f32",
         )
         os.remove("llama_f32.safetensors")