Added hf_llm_engine and config

sia/local_llm_engine.py

@@ -0,0 +1,79 @@
+from threading import Thread
+from typing import Iterator
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
+import torch
+
+from . import util
+from .llm_engine import LlmEngine
+
+class LocalLlmEngine(LlmEngine):
+    def __init__(self, model_path: str):
+        """
+        Initialize the LLM Engine with a model path.
+
+        Args:
+            model_path: Path to the model weights to be used.
+        """
+        self.set_model_path(model_path)
+
+    def set_model_path(self, model_path: str):
+        """
+        Load the model from the specified path.
+
+        Args:
+            model_path: Path to the model weights to load.
+        """
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            return_dict=True,
+            low_cpu_mem_usage=True,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+        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=self.tokenizer,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            return_full_text=False,
+        )
+
+    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
+        """
+        Run inference using the system prompt and main context, while validating actions against the provided XML schema.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_prompt=True
+        )
+        pipeline_kwargs = dict(
+            text_inputs=prompt,
+            do_sample=True,
+            max_new_tokens=1024,
+            streamer=streamer
+        )
+        thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs)
+        thread.start()
+        return util.stop_before_value(streamer, '<|eot_id|>')