27 lines
853 B
Python
27 lines
853 B
Python
import unittest
|
|
|
|
from itertools import tee
|
|
|
|
from . import test_data
|
|
|
|
from sia.llm_engine import LlmEngine
|
|
from sia.local_llm_engine import LocalLlmEngine
|
|
|
|
class LlmEngineTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.model_path = "/root/model"
|
|
|
|
def test_initialization(self):
|
|
llm_engine = LocalLlmEngine(self.model_path)
|
|
self.assertIsInstance(llm_engine, LlmEngine)
|
|
|
|
def test_infer(self):
|
|
main_context = "This is a test"
|
|
llm_engine = LocalLlmEngine(self.model_path)
|
|
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>")
|