V10

app.py

@@ -1,49 +1,112 @@
-# test_local.py
+import os
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+import threading
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
-model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
+# 1) Cargar tokenizer y modelo
+print("Cargando tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Llama-8B")
 
-print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name,
-    trust_remote_code=True  # <-- por si el repo define código extra
-)
-
-print("Loading model...")
+print("Cargando modelo...")
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    trust_remote_code=True,
-    device_map="auto",         # GPU si está disponible
-    torch_dtype=torch.float16   # si GPU; en CPU podrías usar float32
+    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
+    device_map="auto",        # Usa GPU si está disponible
+    torch_dtype=torch.float16  # FP16 si GPU; en CPU, podrías usar float32
 )
-# For safety, set pad_token_id if missing
+model.eval()
+
+# Ajuste de pad_token_id si fuese necesario
 if model.config.pad_token_id is None:
     model.config.pad_token_id = tokenizer.eos_token_id
-model.eval()
 
-print("Testing generation...")
 
-# Prompt sencillo
-prompt = (
-    "System: Eres un asistente amable.\n\n"
-    "User: Hola, ¿cómo estás?\n"
-    "Assistant:"
-)
+def respond(
+    user_message: str,
+    history: list[tuple[str, str]],
+    system_message: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+):
+    """
+    Llamada por ChatInterface en cada turno.
+      - user_message: Texto nuevo del usuario.
+      - history: Lista [(usuario, asistente), ...] de turnos previos.
+      - system_message: Se añade SOLO si el historial está vacío.
+    Devuelve tokens progresivamente (streaming).
+    """
+
+    # 1) Creamos un prompt vacío
+    prompt = ""
+
+    # 2) Solo añadimos system_message si no hay historial
+    if not history:
+        prompt += f"{system_message}\n\n"
 
-inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    # 3) Añadimos la conversación previa: "User: ...\nAssistant: ..."
+    for (past_user, past_assistant) in history:
+        prompt += f"Usuario: {past_user}\nAsistente: {past_assistant}\n"
 
-# Generamos ~50 tokens, do_sample
-with torch.no_grad():
-    output_ids = model.generate(
-        **inputs,
-        max_new_tokens=50,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True
+    # 4) Añadimos el nuevo turno del usuario
+    prompt += f"Usuario: {user_message}\nAsistente:"
+
+    # 5) Preparamos un TextIteratorStreamer para streaming
+    streamer = TextIteratorStreamer(tokenizer=tokenizer, skip_special_tokens=True)
+
+    # 6) Codificamos el prompt
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    # 7) Preparamos parámetros de generate
+    generation_kwargs = {
+        "inputs": inputs["input_ids"],
+        "attention_mask": inputs["attention_mask"],
+        "max_new_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "do_sample": True,
+        "streamer": streamer,
+    }
+
+    # 8) Ejecutamos model.generate en un hilo
+    generation_thread = threading.Thread(
+        target=model.generate,
+        kwargs=generation_kwargs
     )
+    generation_thread.start()
+
+    # 9) Iteramos sobre el streamer para devolver tokens sucesivamente
+    output_text = ""
+    for new_token in streamer:
+        output_text += new_token
+        yield output_text
+
+# 10) Definimos la interfaz ChatInterface
+demo = gr.ChatInterface(
+    fn=respond,
+    additional_inputs=[
+        # Cuadro para "mensaje de sistema", solo usado en la 1ª interacción
+        gr.Textbox(
+            label="Mensaje del sistema (se usará sólo al inicio)",
+            value="Eres un asistente amable y paciente. Responde siempre en español."
+        ),
+        gr.Slider(
+            minimum=1, maximum=1024, value=100, step=1, label="Max new tokens"
+        ),
+        gr.Slider(
+            minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperatura"
+        ),
+        gr.Slider(
+            minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"
+        ),
+    ],
+    title="DeepSeek Chat Demo",
+    description=(
+        "Este demo carga el modelo deepseek-ai/DeepSeek-R1-Distill-Llama-8B "
+        "y permite conversar en varios turnos. El mensaje de sistema se añade "
+        "sólo en la primera interacción, y luego la charla fluye como Usuario/Asistente."
+    )
+)
 
-# Decodificamos y mostramos
-output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-print("\n=== Generated Text ===\n")
-print(output_text)
+if __name__ == "__main__":
+    demo.launch()