from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Carregar o modelo e o tokenizer do GPT-2
model_name = "gpt2"  # Você pode substituir por outros modelos, como EleutherAI/gpt-neo-125M
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def chatbot():
    print("Chatbot: Olá! Sou seu assistente virtual. Digite 'sair' para encerrar a conversa.")
    
    # Histórico de mensagens
    conversation_history = ""

    while True:
        # Entrada do usuário
        user_input = input("Você: ")
        if user_input.lower() in ["sair", "exit", "quit"]:
            print("Chatbot: Até logo!")
            break

        # Adicionar a entrada do usuário ao histórico
        conversation_history += f"Usuário: {user_input}\nChatbot:"

        # Tokenizar a entrada com o histórico
        inputs = tokenizer.encode(conversation_history, return_tensors="pt", truncation=True, max_length=1024)

        # Gerar resposta
        outputs = model.generate(
            inputs,
            max_length=1024,
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id,
            temperature=0.7,
            top_p=0.9
        )

        # Decodificar a resposta gerada
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        chatbot_response = response[len(conversation_history):].strip()

        # Mostrar a resposta
        print(f"Chatbot: {chatbot_response}")

        # Atualizar o histórico de conversas
        conversation_history += f"{chatbot_response}\n"

if __name__ == "__main__":
    chatbot()