Upload src/llm.py with huggingface_hub

src/llm.py

@@ -0,0 +1,30 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class LocalLLM:
+    def __init__(self, model_name: str):
+        from transformers import pipeline
+
+        logger.info("Loading local LLM: %s", model_name)
+        self.pipe = pipeline("text-generation", model=model_name, device=-1)
+        self.tokenizer = self.pipe.tokenizer
+
+    def generate(self, messages: list[dict], max_tokens: int = 1024, temperature: float = 0.1) -> str:
+        prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        result = self.pipe(
+            prompt,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            do_sample=temperature > 0,
+            pad_token_id=self.tokenizer.eos_token_id,
+        )
+        text = result[0]["generated_text"]
+        if text.startswith(prompt):
+            text = text[len(prompt):]
+        return text.strip()
+
+    def close(self):
+        del self.pipe
+        del self.tokenizer