from transformers import AutoModelForCausalLM, AutoTokenizer
from config.model_config import HIMConfig
import torch

class HIMModel:
    def __init__(self, config: HIMConfig):
        self.config = config
        self.tokenizer = AutoTokenizer.from_pretrained(config.base_model)
        self.model = AutoModelForCausalLM.from_pretrained(config.base_model)
        
    def generate_response(self, input_text: str, system_message: str = None):
        # Prepare input with system message if provided
        if system_message:
            input_text = f"{system_message}\nUser: {input_text}\nHIM:"
            
        inputs = self.tokenizer(input_text, return_tensors="pt")
        
        outputs = self.model.generate(
            inputs["input_ids"],
            max_length=self.config.max_length,
            temperature=self.config.temperature,
            top_p=self.config.top_p,
            do_sample=True
        )
        
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)