Merge pull request agiresearch#25 from dongyuanjushi/main

add test code and dependencies via pytest

requirements.txt

@@ -15,3 +15,4 @@ torch==2.0.1
 transformers==4.38.1
 wolframalpha==5.0.0
 pre-commit==3.7.0
+pytest==8.1.1

src/llms/llms.py

@@ -24,7 +24,7 @@
 import time
 
 class LLMKernel:
-    def __init__(self, llm_name: str, max_gpu_memory: dict, eval_device: str, max_new_tokens: int = 256):
+    def __init__(self, llm_name: str, max_gpu_memory: dict = None, eval_device: str = None, max_new_tokens: int = 256):
         print("Initialize AIOS powered by LLM: {}".format(llm_name))
         self.config = self.load_config(llm_name)
         self.max_gpu_memory = max_gpu_memory

tests/test_llms.py

@@ -0,0 +1,13 @@
+import sys
+
+import os
+
+from src.llms.llms import LLMKernel
+
+def test_closed_llm():
+    llm_type = "gemini-pro"
+    llm = LLMKernel(llm_type)
+    prompt = "Craft a tale about a valiant warrior on a quest to uncover priceless treasures hidden within a mystical island."
+    response = llm.address_request(prompt)
+    assert isinstance(response, str)
+    print(response)