66 lines
2.3 KiB
Python
66 lines
2.3 KiB
Python
import unittest
|
|
from sia.llm_engine import LlmEngine, InferenceResult
|
|
|
|
echo_system_prompt = """
|
|
Your answer always consists of 4 parts:
|
|
- The original request
|
|
- <test_tag>
|
|
- The original request
|
|
- </test_tag>
|
|
You never provide an answer.
|
|
Only the entered text, the xml open tag, the same text again and the xml close tag.
|
|
Don't add whitespace or newlines.
|
|
Don't modify the text or change casing.
|
|
Be exact, this is for testing purposes.
|
|
"""
|
|
|
|
echo_action_schema = """
|
|
<?xml version="1.0" encoding="UTF-8"?>
|
|
<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
|
|
<xs:element name="test_tag">
|
|
<xs:complexType>
|
|
<xs:simpleContent>
|
|
<xs:extension base="xs:string">
|
|
<xs:attribute name="id" type="xs:string" use="required"/>
|
|
</xs:extension>
|
|
</xs:simpleContent>
|
|
</xs:complexType>
|
|
</xs:element>
|
|
</xs:schema>
|
|
"""
|
|
|
|
class TestLlmEngine(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)
|
|
result = llm_engine.infer(echo_system_prompt, main_context, 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>")
|
|
|
|
'''
|
|
def test_set_model_path(self):
|
|
new_model_path = "/path/to/new/model"
|
|
self.llm_engine.set_model_path(new_model_path)
|
|
# Add assertions to check if the model path was updated correctly
|
|
|
|
def test_finetune(self):
|
|
dataset_paths = ["/path/to/dataset1", "/path/to/dataset2"]
|
|
output_dir = "/path/to/output"
|
|
|
|
self.llm_engine.finetune(dataset_paths, output_dir)
|
|
# Add assertions to check if the fine-tuning process completed successfully
|
|
# For example, check if new model weights were saved in the output directory
|
|
'''
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main() |