from transformers import AutoModelForCausalLM, AutoTokenizer

class ChatbotModel:
    def __init__(self, model_name: str):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(model_name)

    def get_response(self, user_message: str) -> str:
        # Tokenize input and generate response
        inputs = self.tokenizer.encode(user_message + self.tokenizer.eos_token, return_tensors="pt")
        outputs = self.model.generate(inputs, max_length=150, num_return_sequences=1)
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response.strip()