Added hf_llm_engine and config

This commit is contained in:
Niels Geens
2024-11-04 17:08:52 +01:00
parent 5da6dca5ec
commit 70ed16f8ab
12 changed files with 330 additions and 93 deletions

View File

@@ -1,25 +0,0 @@
import unittest
from itertools import tee
from . import test_data
from sia.llm_engine import LlmEngine
class LlmEngineTest(unittest.TestCase):
def setUp(self):
self.model_path = "/root/model"
def test_initialization(self):
llm_engine = LlmEngine(self.model_path)
self.assertIsInstance(llm_engine, LlmEngine)
def test_infer(self):
main_context = "This is a test"
llm_engine = LlmEngine(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>")