app.py

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

# Cargar el modelo y el tokenizador
model_name = "bigscience/bloom-560m"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def chatbot(input, history):
    history = history or []
    history.append({"role": "user", "content": input})
    
    chat_history = ""
    for message in history:
        if message["role"] == "user":
            chat_history += f"Human: {message['content']}\n"
        else:
            chat_history += f"AI: {message['content']}\n"
    
    chat_history += "AI:"
    
    input_ids = tokenizer.encode(chat_history, return_tensors="pt")
    attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device)
    
    max_length = input_ids.shape[1] + 50
    
    output = model.generate(
        input_ids, 
        attention_mask=attention_mask,
        max_length=max_length,
        num_return_sequences=1,
        no_repeat_ngram_size=2,
        temperature=0.7
    )
    
    response = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
    
    history.append({"role": "assistant", "content": response.strip()})
    
    # Convertir el historial al formato que Gradio espera
    gradio_history = [[m["content"], h["content"]] for m, h in zip(history[::2], history[1::2])]
    
    return gradio_history, history

iface = gr.Interface(
    fn=chatbot,
    inputs=["text", "state"],
    outputs=["chatbot", "state"],
    title="Tu Compañero AI con BLOOM",
    description="Un chatbot de IA diseñado para simular conversaciones personales, utilizando el modelo BLOOM.",
)

iface.launch()