llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Start work on docker module

Browse Source
This commit is contained in:
geens
2024-10-24 08:39:31 +02:00
parent ab03e6836b
commit 4bbe636f0e
8 changed files with 532 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
test/llm_engine_test.py
View File

@@ -1,14 +1,15 @@
import unittest
from itertools import tee
from . import test_data
from . import test_util
from sia.llm_engine import LlmEngine, InferenceResult
from sia.llm_engine import LlmEngine
class LlmEngineTest(unittest.TestCase):
def setUp(self):
self.model_path = "/root/model"
self.llm_engine = LlmEngine(self.model_path)
def test_initialization(self):
llm_engine = LlmEngine(self.model_path)
@@ -17,9 +18,9 @@ class LlmEngineTest(unittest.TestCase):
def test_infer(self):
main_context = "This is a test"
llm_engine = LlmEngine(self.model_path)
result = llm_engine.infer(test_data.echo_system_prompt, main_context, test_data.echo_action_schema)
self.assertIsInstance(result, InferenceResult)
self.assertIsInstance(result.reasoning, str)
self.assertIsInstance(result.actions, str)
self.assertEqual(result.reasoning, main_context)
self.assertEqual(result.actions, f"<test_tag>{main_context}</test_tag>")
tokens = llm_engine.infer(test_data.echo_system_prompt, main_context)
print_tokens, result_tokens = tee(tokens)
for token in print_tokens:
print(token, end="", flush=True)
result = ''.join(result_tokens)
self.assertEqual(result, f"{main_context}<test_tag>{main_context}</test_tag>")