Begin implementation, basic inferenece

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):
         """