Update app.py

app.py

@@ -1,47 +1,54 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 # Cargar el modelo y el tokenizador
-model_name = "microsoft/DialoGPT-medium"
+model_name = "bigscience/bloom-560m"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
 def chatbot(input, history):
-    # Asegúrate de que history sea una lista de listas
     history = history or []
+    history.append({"role": "user", "content": input})
     
-    # Agregar el input del usuario al historial
-    history.append([input, None])
+    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"
     
-    # Preparar el contexto para el modelo
-    chat_history = []
-    for human, ai in history:
-        chat_history.append(human)
-        if ai:
-            chat_history.append(ai)
+    chat_history += "AI:"
     
-    # Tokenizar la conversación
-    input_ids = tokenizer.encode(" ".join(chat_history), return_tensors="pt")
+    input_ids = tokenizer.encode(chat_history, return_tensors="pt")
+    attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device)
     
-    # Generar una respuesta
-    attention_mask = input_ids.new_ones(input_ids.shape)
-    output = model.generate(input_ids, attention_mask=attention_mask, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+    max_length = input_ids.shape[1] + 50
     
-    # Decodificar la respuesta
-    response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+    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
+    )
     
-    # Agregar la respuesta al historial
-    history[-1][1] = response
+    response = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
     
-    return history, history
+    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
 
-# Crear la interfaz de Gradio
 iface = gr.Interface(
     fn=chatbot,
     inputs=["text", "state"],
     outputs=["chatbot", "state"],
-    title="Tu Compañero AI",
-    description="Un chatbot de IA diseñado para simular conversaciones personales.",
+    title="Tu Compañero AI con BLOOM",
+    description="Un chatbot de IA diseñado para simular conversaciones personales, utilizando el modelo BLOOM.",
 )
 
 iface.launch()