Begin implementation, basic inferenece

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)