Begin implementation, basic inferenece

Dockerfile

@@ -5,8 +5,9 @@ COPY ./sia/ /root/sia/
 RUN pip3 install -r requirements.txt
 
 FROM requirements AS test
-COPY ./tests/ /root/tests/
+COPY ./test/ /root/test/
 RUN mkdir -p /root/model
-CMD python3 -m unittest discover tests
+CMD ["python3", "-m", "unittest", "discover", "-p", "*test.py", "-v"]
 
-FROM requirements
+FROM requirements
+CMD ["python3", "-m", "sia"]

sia/__main__.py

@@ -0,0 +1,5 @@
+def main():
+    print("Hello, World! --sia")
+
+if __name__ == "__main__":
+    main()

sia/inference_result.py

@@ -0,0 +1,5 @@
+from typing import NamedTuple
+
+class InferenceResult(NamedTuple):
+    reasoning: str
+    actions: str

sia/llm_engine.py

@@ -1,10 +1,8 @@
-from typing import NamedTuple
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-class InferenceResult(NamedTuple):
-    reasoning: str
-    actions: str
+from . import util
+from .inference_result import InferenceResult
 
 class LlmEngine:
     def __init__(self, model_path: str):
@@ -14,9 +12,6 @@ class LlmEngine:
         Args:
             model_path: Path to the model weights to be used.
         """
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        print(f"device: {self.device}")
         self.set_model_path(model_path)
 
     def set_model_path(self, model_path: str):
@@ -26,24 +21,24 @@ class LlmEngine:
         Args:
             model_path: Path to the model weights to load.
         """
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
         model = AutoModelForCausalLM.from_pretrained(
             model_path,
             return_dict=True,
             low_cpu_mem_usage=True,
-            torch_dtype=self.torch_dtype,
+            torch_dtype=torch.bfloat16,
             device_map="auto",
             trust_remote_code=True,
-        ).to(self.device)
-        if tokenizer.pad_token_id is None:
-            tokenizer.pad_token_id = tokenizer.eos_token_id
+        )
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
         if model.config.pad_token_id is None:
             model.config.pad_token_id = model.config.eos_token_id
         self.pipeline = pipeline(
             "text-generation",
             model=model,
-            tokenizer=tokenizer,
-            torch_dtype=torch.float16,
+            tokenizer=self.tokenizer,
+            torch_dtype=torch.bfloat16,
             device_map="auto",
         )
 
@@ -57,9 +52,21 @@ class LlmEngine:
             action_schema: XML schema to validate the generated actions
 
         Returns:
-            InferenceResult: the actions validate against the schema
+            InferenceResult: Tuple containing reasoning and actions that validate against the schema
         """
-        pass
+        valid_elements = util.get_valid_root_elements(action_schema)
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        outputs = self.pipeline(prompt, max_new_tokens=120, do_sample=True)
+        generated_text = outputs[0]["generated_text"]
+        response = generated_text.split("<|start_header_id|>assistant<|end_header_id|>",1)[1].strip()
+        result = util.split_response(response, valid_elements)
+        return result
 
     def finetune(self, dataset_paths: list, output_dir: str):
         """

sia/util.py

@@ -0,0 +1,47 @@
+from .inference_result import InferenceResult
+
+import xml.etree.ElementTree as ET
+import re
+
+def get_valid_root_elements(schema: str) -> set:
+    """
+    Extract valid root element names from the XML schema.
+    
+    Args:
+        schema: XML schema string
+        
+    Returns:
+        set: Set of valid root element names
+    """
+    try:
+        schema = schema.strip()
+        ET.register_namespace('xs', 'http://www.w3.org/2001/XMLSchema')
+        root = ET.fromstring(schema)
+        ns = {'xs': 'http://www.w3.org/2001/XMLSchema'}
+        elements = root.findall(".//xs:element", ns)
+        return {elem.get('name') for elem in elements if elem.get('name')}
+    except ET.ParseError as e:
+        print(f"Error parsing schema: {e}")
+        return set()
+
+def split_response(response: str, valid_elements: set) -> InferenceResult:
+    """
+    Split the response into reasoning and actions based on valid XML elements.
+    
+    Args:
+        response: Raw response string from the model
+        valid_elements: Set of valid root element names from the schema
+        
+    Returns:
+        InferenceResult: Tuple containing reasoning and actions
+    """
+    elements_pattern = '|'.join(map(re.escape, valid_elements))
+    pattern = f"<({elements_pattern})[^>]*>"
+    matches = list(re.finditer(pattern, response))
+    if not matches:
+        return InferenceResult(response.strip(), "")
+    last_match = matches[-1]
+    split_point = last_match.start()
+    reasoning = response[:split_point].strip()
+    actions = response[split_point:].strip()
+    return InferenceResult(reasoning, actions)

test/llm_engine_test.py

@@ -0,0 +1,25 @@
+import unittest
+
+from . import test_data
+from . import test_util
+
+from sia.llm_engine import LlmEngine, InferenceResult
+
+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)
+        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(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>")

test/test_data.py

@@ -0,0 +1,32 @@
+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.
+
+    example input:
+    hello
+    example output:
+    hello<test_tag>hello</test_tag>
+""".strip()
+
+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>
+""".strip()

test/test_util.py

@@ -0,0 +1,8 @@
+import unittest
+
+class SequentialTestLoader(unittest.TestLoader):
+    def getTestCaseNames(self, testCaseClass):
+        test_names = super().getTestCaseNames(testCaseClass)
+        testcase_methods = list(testCaseClass.__dict__.keys())
+        test_names.sort(key=testcase_methods.index)
+        return test_names

test/util_test.py

@@ -0,0 +1,17 @@
+import unittest
+
+from . import test_data
+
+from sia import util
+
+class UtilTest(unittest.TestCase):
+    def test_get_valid_root_elements_single(self):
+        valid_elements = util.get_valid_root_elements(test_data.echo_action_schema)
+        self.assertEqual(valid_elements, {'test_tag'})
+
+    def test_split_response_single_element(self):
+        response = "Some reasoning here\n<test_tag>content</test_tag>"
+        valid_elements = {'test_tag'}
+        result = util.split_response(response, valid_elements)
+        self.assertEqual(result.reasoning, "Some reasoning here")
+        self.assertEqual(result.actions, "<test_tag>content</test_tag>")

tests/test_llm_engine.py

@@ -1,66 +0,0 @@
-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()