src/llm.py

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