From 05aab43c312434f23160e91ad05629b5939a1c37 Mon Sep 17 00:00:00 2001 From: saienduri Date: Mon, 27 Nov 2023 10:09:23 -0800 Subject: [PATCH] update test --- python/turbine_models/tests/llama_test.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/python/turbine_models/tests/llama_test.py b/python/turbine_models/tests/llama_test.py index 903f626d6..1b273c05a 100644 --- 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")