53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
import unittest
|
|
|
|
from itertools import tee
|
|
|
|
from . import test_data
|
|
|
|
from sia.llm_engine import LlmEngine
|
|
from sia.llm_engine.local_llm_engine import LocalLlmEngine
|
|
|
|
class LocalLlmEngineTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.model_path = "/root/models/current"
|
|
|
|
@unittest.skip("Takes too long")
|
|
def test_initialization(self):
|
|
llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
|
|
self.assertIsInstance(llm_engine, LlmEngine)
|
|
|
|
@unittest.skip("Takes too long")
|
|
def test_infer(self):
|
|
main_context = "This is a test"
|
|
llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
|
|
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>")
|
|
|
|
@unittest.skip("Takes too long")
|
|
def test_token_count(self):
|
|
"""Test token counting returns reasonable numbers"""
|
|
llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
|
|
|
|
# Test with short inputs
|
|
count = llm_engine.token_count("Short prompt", "Brief context")
|
|
self.assertGreater(count, 5)
|
|
self.assertLess(count, 50)
|
|
|
|
# Test with longer inputs
|
|
long_prompt = "A detailed system prompt with multiple sentences. " * 5
|
|
long_context = "An extensive context containing various details. " * 10
|
|
count = llm_engine.token_count(long_prompt, long_context)
|
|
self.assertGreater(count, 100)
|
|
self.assertLess(count, 1000)
|
|
|
|
@unittest.skip("Takes too long")
|
|
def test_token_limit(self):
|
|
"""Test token limit is within expected range for modern LLMs"""
|
|
llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
|
|
limit = llm_engine.token_limit()
|
|
self.assertGreaterEqual(limit, 2048)
|
|
self.assertLessEqual(limit, 1048576) |