Switch back to 7b

nod-ai · Jan 17, 2024 · 1c7896d · 1c7896d
1 parent 1ed5543
commit 1c7896d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/apps/shark_studio/tests/api_test.py b/apps/shark_studio/tests/api_test.py
@@ -13,12 +13,12 @@
 class LLMAPITest(unittest.TestCase):
     def test01_LLMSmall(self):
         lm = LanguageModel(
-            "TinyPixel/small-llama2",
+            "Trelis/Llama-2-7b-chat-hf-function-calling-v2",
             hf_auth_token=None,
             device="cpu",
             external_weights="safetensors",
             precision="fp32",
-            quantization="None",
+            quantization="int4",
         )
         count = 0
         for msg, _ in lm.chat("hi, what are you?"):
@@ -27,8 +27,8 @@ def test01_LLMSmall(self):
                 count += 1
                 continue
             assert (
-                msg.strip(" ") == "Turkish Turkish Turkish"
-            ), f"LLM API failed to return correct response, expected 'Turkish Turkish Turkish', received {msg}"
+                msg.strip(" ") == "Hello!"
+            ), f"LLM API failed to return correct response, expected 'Hello!', received {msg}"
             break
         del lm
         gc.collect()