import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# Carga el modelo y el tokenizer
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def predict(input_text, history=[]):
    # Tokeniza la nueva entrada del usuario
    new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")

    # Concatena la nueva entrada del usuario con el historial
    bot_input_ids = torch.cat([torch.tensor(history, dtype=torch.long), new_user_input_ids], dim=-1)

    # Genera una respuesta
    output_ids = model.generate(bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id).tolist()[0]

    # Extrae solo la parte de la respuesta
    response = tokenizer.decode(output_ids, skip_special_tokens=True)

    # Actualiza el historial
    history = output_ids if len(output_ids) < 4000 else output_ids[-4000:]

    return response, history

# Configuración de Gradio Interface
title = "Conversational AI Chatbot"
description = "Este es un chatbot que utiliza un modelo de lenguaje para generar respuestas. Introduce tu mensaje y recibe una respuesta."
examples = [["Hello, how are you?"], ["What is your name?"]]

iface = gr.Interface(
    fn=predict,
    title=title,
    description=description,
    examples=examples,
    inputs=[gr.Textbox(label="Tu Mensaje"), gr.State()],
    outputs=[gr.Textbox(label="Respuesta del Chatbot"), gr.State()],
    allow_flagging="never",
).launch()