fix: llm test

dvilelaf · May 30, 2024 · ce9ab3a · ce9ab3a
1 parent 5783a44
commit ce9ab3a
Showing 1 changed file with 41 additions and 19 deletions.
diff --git a/scripts/test_prompts.py b/scripts/test_prompts.py
@@ -20,6 +20,8 @@
 
 """Script to test agent prompts"""
 
+import time
+
 from llama_cpp import Llama
 
 from packages.dvilela.skills.tsunami_abci.prompts import (
@@ -28,28 +30,48 @@
 )
 
 
+def timer(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        print(f"Elapsed time for {func.__name__}: {elapsed_time} seconds")
+        return result
+
+    return wrapper
+
+
 llm = Llama.from_pretrained(
     repo_id="TheBloke/CapybaraHermes-2.5-Mistral-7B-GGUF",
-    filename="*Q2_K.gguf",
+    # repo_id="QuantFactory/Meta-Llama-3-8B-Instruct-GGUF",
+    # filename="*Q2_K.gguf",
+    filename="*Q8_0.gguf",
     verbose=False,
 )
 
-output = llm.create_chat_completion(
-    messages=[
-        {"role": "system", "content": SYSTEM_PROMPTS[-1]},  # type: ignore
-        {  # type: ignore
-            "role": "user",
-            "content": OMEN_USER_PROMPT.format(
-                n_markets=15,
-                n_agents=20,
-                n_trades=900,
-                usd_amount=100,
-                biggest_trader_address="0x0000000000",
-                biggest_trader_trades=30,
-            ),
-        },
-    ],
-    temperature=0.8,
-)
 
-print(output["choices"][0]["message"]["content"])  # type: ignore
+@timer
+def run_llm():
+    output = llm.create_chat_completion(
+        messages=[
+            {"role": "system", "content": SYSTEM_PROMPTS[-1]},  # type: ignore
+            {  # type: ignore
+                "role": "user",
+                "content": OMEN_USER_PROMPT.format(
+                    n_markets=15,
+                    n_agents=20,
+                    n_trades=900,
+                    usd_amount=100,
+                    biggest_trader_address="0x0000000000",
+                    biggest_trader_trades=30,
+                ),
+            },
+        ],
+        temperature=0.8,
+    )
+
+    print(output["choices"][0]["message"]["content"])  # type: ignore
+
+
+run_llm()