from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

class LlamaChatbot:
    def __init__(self, model_name="meta-llama/Llama-3.2-7b"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(model_name)
    
    def get_response(self, user_input: str, conversation_history: str) -> str:
        # Combine user input with the conversation history
        input_text = conversation_history + user_input + " "
        
        inputs = self.tokenizer(input_text, return_tensors="pt")
        outputs = self.model.generate(**inputs, max_length=150)
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return response.strip()