Update app.py

app.py

@@ -1,18 +1,24 @@
-# app.py — DeepSeek-OCR (GPU worker) + TxAgent-T1-Llama-3.1-8B (HF Inference vía text_generation)
-# -----------------------------------------------------------------------------------------------
-# • OCR: DeepSeek-OCR cargado en CPU y movido a GPU SOLO dentro de @spaces.GPU (evita “CUDA en main”).
-# • Chat: mims-harvard/TxAgent-T1-Llama-3.1-8B por InferenceClient.text_generation con provider Featherless AI.
-# • Sin queue(concurrency_count): compatible con Gradio 5.
-# • Variables recomendadas (Settings → Secrets):
-#     HF_TOKEN=hf_xxx                      (requerido para Inference)
-#     TX_MODEL_ID=mims-harvard/TxAgent-T1-Llama-3.1-8B
-#     TX_TOKENIZER_ID=mims-harvard/TxAgent-T1-Llama-3.1-8B
-#     GEN_MAX_NEW_TOKENS=512
-#     GEN_TEMPERATURE=0.2
-#     GEN_TOP_P=0.9
-#     OCR_REVISION=<commit opcional para fijar versión estable>
-#     OCR_ATTN_IMPL=flash_attention_2   (o "eager" si no hay FlashAttention)
-# -----------------------------------------------------------------------------------------------
+# app.py — DeepSeek-OCR (GPU worker) + TxAgent-T1-Llama-3.1-8B (HF Inference conversational)
+# ------------------------------------------------------------------------------------------------
+# Flujo:
+#   1. OCR local con DeepSeek-OCR (CUDA solo dentro de @spaces.GPU).
+#   2. Chat médico remoto con TxAgent-T1-Llama-3.1-8B usando provider "featherless-ai"
+#      vía .chat.completions.create() (tarea conversational).
+#
+# Variables de entorno recomendadas (Settings → Secrets):
+#   HF_TOKEN=hf_xxx                 (OBLIGATORIO para usar inference)
+#   TX_MODEL_ID=mims-harvard/TxAgent-T1-Llama-3.1-8B
+#   GEN_MAX_NEW_TOKENS=512
+#   GEN_TEMPERATURE=0.2
+#   GEN_TOP_P=0.9
+#   OCR_REVISION=<commit opcional estable de DeepSeek-OCR>
+#   OCR_ATTN_IMPL=flash_attention_2   (o "eager" si no hay FlashAttention2)
+#
+# Nota importante:
+#   - NO tocamos CUDA en el proceso principal. Solo dentro de ocr_infer().
+#   - No usamos text_generation. El provider featherless-ai ofrece "conversational".
+#   - Evitamos el 404 del router hf-inference porque forzamos provider="featherless-ai".
+# ------------------------------------------------------------------------------------------------
 
 import os, tempfile, traceback
 import gradio as gr
@@ -23,96 +29,84 @@ import spaces
 from huggingface_hub import InferenceClient
 
 # =========================
-# Config — Chat remoto (TxAgent por text_generation + Featherless)
+# Configuración del Chat remoto (TxAgent)
 # =========================
-TX_MODEL_ID      = os.getenv("TX_MODEL_ID", "mims-harvard/TxAgent-T1-Llama-3.1-8B")
-TX_TOKENIZER_ID  = os.getenv("TX_TOKENIZER_ID", TX_MODEL_ID)
-HF_TOKEN         = os.getenv("HF_TOKEN")  # requerido
+TX_MODEL_ID = os.getenv("TX_MODEL_ID", "mims-harvard/TxAgent-T1-Llama-3.1-8B")
+HF_TOKEN    = os.getenv("HF_TOKEN")
 
 GEN_MAX_NEW_TOKENS = int(os.getenv("GEN_MAX_NEW_TOKENS", "512"))
 GEN_TEMPERATURE    = float(os.getenv("GEN_TEMPERATURE", "0.2"))
 GEN_TOP_P          = float(os.getenv("GEN_TOP_P", "0.9"))
 
-# Cliente genérico (sin atar modelo/proveedor; se pasa en cada llamada)
-_hf_client = InferenceClient(token=HF_TOKEN, timeout=60.0)
-
-# Tokenizer para aplicar chat template → prompt
-_TX_TOKENIZER = None
-def get_tx_tokenizer():
-    global _TX_TOKENIZER
-    if _TX_TOKENIZER is None:
-        _TX_TOKENIZER = AutoTokenizer.from_pretrained(TX_TOKENIZER_ID, trust_remote_code=True)
-    return _TX_TOKENIZER
+# Cliente remoto del modelo.
+# Clave: provider="featherless-ai", que es el que sí soporta este modelo en modo conversational.
+tx_client = InferenceClient(
+    model=TX_MODEL_ID,
+    provider="featherless-ai",
+    token=HF_TOKEN,
+    timeout=60.0,
+)
 
 def _system_prompt():
     return (
         "Eres un asistente clínico educativo. NO sustituyes el juicio médico.\n"
-        "Usa CONTEXTO_OCR si existe; si falta, dilo explícitamente. No inventes datos fuera del OCR."
+        "Usa CONTEXTO_OCR si existe; si falta, dilo explícitamente. "
+        "No inventes datos que no estén en el OCR ni hagas diagnósticos definitivos."
     )
 
-def _mk_messages(ocr_md: str, ocr_txt: str, user_msg: str):
+def _mk_messages_for_provider(ocr_md: str, ocr_txt: str, user_msg: str):
+    """
+    Este formato es exactamente el que espera chat.completions.create():
+    lista de dicts con role: system/user/assistant.
+    """
     ctx = (ocr_md or "")[:3000] or (ocr_txt or "")[:3000]
-    sys = _system_prompt()
+
+    sys_content = _system_prompt()
     if ctx:
-        sys += "\n\n---\nCONTEXTO_OCR (fuente principal):\n" + ctx + "\n---"
+        sys_content += (
+            "\n\n---\n"
+            "CONTEXTO_OCR (extraído de la imagen):\n"
+            f"{ctx}\n"
+            "---\n"
+            "Responde basándote en ese contenido. Si falta información, dilo."
+        )
+
     if not user_msg:
-        user_msg = "Analiza el CONTEXTO_OCR anterior y responde a partir de ese contenido."
+        user_msg = (
+            "Analiza el CONTEXTO_OCR anterior y explícame, en lenguaje claro, "
+            "qué medicamentos aparecen, dosis y advertencias importantes."
+        )
+
     return [
-        {"role": "system", "content": sys},
+        {"role": "system", "content": sys_content},
         {"role": "user", "content": user_msg},
     ]
 
 def txagent_chat_remote(ocr_md: str, ocr_txt: str, user_msg: str) -> str:
     """
-    Usa text_generation con provider Featherless AI.
-    - Convertimos mensajes a prompt con el chat template del tokenizer.
-    - Llamamos al router con model='mims-harvard/TxAgent…:featherless-ai'
+    Llama a la tarea 'conversational' del provider featherless-ai para TxAgent.
+    Esto evita el error:
+      - text-generation no soportado
+      - 404 de hf-inference
     """
-    messages = _mk_messages(ocr_md, ocr_txt, user_msg)
-    tok = get_tx_tokenizer()
-    prompt = tok.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True,  # deja el turno del assistant abierto
-    )
+    messages = _mk_messages_for_provider(ocr_md, ocr_txt, user_msg)
 
-    model_with_provider = f"{TX_MODEL_ID}:featherless-ai"
     try:
-        out = _hf_client.text_generation(
-            model=model_with_provider,
-            prompt=prompt,
-            max_new_tokens=GEN_MAX_NEW_TOKENS,
+        completion = tx_client.chat.completions.create(
+            model=TX_MODEL_ID,
+            messages=messages,
+            max_tokens=GEN_MAX_NEW_TOKENS,
             temperature=GEN_TEMPERATURE,
             top_p=GEN_TOP_P,
             stream=False,
         )
-        # En huggingface_hub nuevas, text_generation devuelve str (texto generado).
-        return out if isinstance(out, str) else str(out)
-    except Exception as e1:
-        # Fallback: crear cliente amarrado al provider por si el mapping cambia
-        try:
-            client_fb = InferenceClient(
-                model=TX_MODEL_ID,
-                provider="featherless-ai",
-                token=HF_TOKEN,
-                timeout=60.0,
-            )
-            out = client_fb.text_generation(
-                prompt=prompt,
-                max_new_tokens=GEN_MAX_NEW_TOKENS,
-                temperature=GEN_TEMPERATURE,
-                top_p=GEN_TOP_P,
-                stream=False,
-            )
-            return out if isinstance(out, str) else str(out)
-        except Exception as e2:
-            raise RuntimeError(
-                f"Remote generation failed: {e1.__class__.__name__}: {e1} | "
-                f"Fallback: {e2.__class__.__name__}: {e2}"
-            )
+        # El objeto completion tiene .choices[i].message.content
+        return completion.choices[0].message.content
+    except Exception as e:
+        raise RuntimeError(f"Inference error: {e.__class__.__name__}: {e}")
 
 # =========================
-# OCR — DeepSeek-OCR (Transformers), CUDA solo en worker GPU
+# OCR local — DeepSeek-OCR
 # =========================
 def _best_dtype():
     if torch.cuda.is_available():
@@ -120,59 +114,81 @@ def _best_dtype():
     return torch.float32
 
 def _load_ocr_model():
+    """
+    Cargamos DeepSeek-OCR con trust_remote_code.
+    IMPORTANTE: no movemos a CUDA aquí. Eso solo ocurre en el worker @spaces.GPU.
+    """
     model_id = "deepseek-ai/DeepSeek-OCR"
-    revision = os.getenv("OCR_REVISION", None)  # fija un commit para estabilidad si quieres
+    revision = os.getenv("OCR_REVISION", None)  # pin commit para evitar que cambie el repo
     attn_impl = os.getenv("OCR_ATTN_IMPL", "flash_attention_2")
 
-    tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, revision=revision)
+    ocr_tokenizer = AutoTokenizer.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+        revision=revision,
+    )
     try:
-        mdl = AutoModel.from_pretrained(
+        ocr_model = AutoModel.from_pretrained(
             model_id,
             trust_remote_code=True,
             use_safetensors=True,
             _attn_implementation=attn_impl,
             revision=revision,
         ).eval()
-        return tok, mdl
+        return ocr_tokenizer, ocr_model
     except Exception as e:
-        # Fallback si FA2 no está disponible en el entorno
+        # fallback sin FlashAttention2
         if any(k in str(e).lower() for k in ["flash_attn", "flashattention2", "flash_attention_2"]):
-            mdl = AutoModel.from_pretrained(
+            ocr_model = AutoModel.from_pretrained(
                 model_id,
                 trust_remote_code=True,
                 use_safetensors=True,
                 _attn_implementation="eager",
                 revision=revision,
             ).eval()
-            return tok, mdl
+            return ocr_tokenizer, ocr_model
         raise
 
 OCR_TOKENIZER, OCR_MODEL = _load_ocr_model()
 
-@spaces.GPU  # ← toca CUDA solo aquí, no en el proceso principal
+@spaces.GPU  # <- ÚNICO sitio donde tocamos CUDA. Cumple con la política de Spaces Zero.
 def ocr_infer(image: Image.Image, model_size: str, task_type: str, is_eval_mode: bool):
+    """
+    Ejecuta OCR en GPU (si hay) y devuelve:
+      - imagen anotada (puede ser None en eval_mode)
+      - markdown OCR
+      - texto llano OCR
+    """
     if image is None:
         return None, "Sube una imagen primero.", "Sube una imagen primero."
 
     dtype = _best_dtype()
-    model = OCR_MODEL.cuda().to(dtype) if torch.cuda.is_available() else OCR_MODEL.to(dtype)
+    model_local = OCR_MODEL.cuda().to(dtype) if torch.cuda.is_available() else OCR_MODEL.to(dtype)
 
     with tempfile.TemporaryDirectory() as outdir:
-        prompt = "<image>\nFree OCR. " if task_type == "Free OCR" else "<image>\n<|grounding|>Convert the document to markdown. "
+        # prompt según modo
+        if task_type == "Free OCR":
+            prompt = "<image>\nFree OCR. "
+        else:
+            prompt = "<image>\n<|grounding|>Convert the document to markdown. "
 
         size_cfgs = {
             "Tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},
             "Small": {"base_size": 640, "image_size": 640, "crop_mode": False},
             "Base": {"base_size": 1024, "image_size": 1024, "crop_mode": False},
             "Large": {"base_size": 1280, "image_size": 1280, "crop_mode": False},
-            "Gundam (Recommended)": {"base_size": 1024, "image_size": 640, "crop_mode": True},
+            "Gundam (Recommended)": {
+                "base_size": 1024,
+                "image_size": 640,
+                "crop_mode": True
+            },
         }
         cfg = size_cfgs.get(model_size, size_cfgs["Gundam (Recommended)"])
 
         tmp_path = os.path.join(outdir, "tmp.jpg")
         image.save(tmp_path)
 
-        plain = model.infer(
+        plain_text_result = model_local.infer(
             OCR_TOKENIZER,
             prompt=prompt,
             image_file=tmp_path,
@@ -185,27 +201,44 @@ def ocr_infer(image: Image.Image, model_size: str, task_type: str, is_eval_mode:
             eval_mode=is_eval_mode,
         )
 
-        img_boxes = os.path.join(outdir, "result_with_boxes.jpg")
-        md_path   = os.path.join(outdir, "result.mmd")
+        img_boxes_path = os.path.join(outdir, "result_with_boxes.jpg")
+        md_path        = os.path.join(outdir, "result.mmd")
 
-        md = "Markdown result was not generated. This is expected for 'Free OCR' task."
+        markdown_content = (
+            "Markdown result was not generated. This is expected for 'Free OCR' task."
+        )
         if os.path.exists(md_path):
             with open(md_path, "r", encoding="utf-8") as f:
-                md = f.read()
+                markdown_content = f.read()
+
+        annotated_img = None
+        if os.path.exists(img_boxes_path):
+            annotated_img = Image.open(img_boxes_path)
+            annotated_img.load()
 
-        img_out = Image.open(img_boxes) if os.path.exists(img_boxes) else None
-        txt_out = plain if plain else md
-        return img_out, md, txt_out
+        text_out = plain_text_result if plain_text_result else markdown_content
+        return annotated_img, markdown_content, text_out
 
 # =========================
-# Glue OCR→Chat
+# Estados / helpers para la UI
 # =========================
 def ocr_snapshot(md_text: str, plain_text: str):
+    """
+    Guardamos el OCR en estados (para enviarlo al chat después)
+    y devolvemos esas vistas rápidas.
+    """
     return md_text, plain_text, md_text, plain_text
 
 def chat_reply(user_msg, chat_state, ocr_md_state, ocr_txt_state):
+    """
+    Lógica del botón "Enviar" en el chat.
+    """
     try:
-        answer = txagent_chat_remote(ocr_md_state or "", ocr_txt_state or "", user_msg or "")
+        answer = txagent_chat_remote(
+            ocr_md_state or "",
+            ocr_txt_state or "",
+            user_msg or ""
+        )
         updated = (chat_state or []) + [
             {"role": "user", "content": user_msg or "(solo OCR)"},
             {"role": "assistant", "content": answer},
@@ -215,7 +248,10 @@ def chat_reply(user_msg, chat_state, ocr_md_state, ocr_txt_state):
         tb = traceback.format_exc(limit=2)
         updated = (chat_state or []) + [
             {"role": "user", "content": user_msg or ""},
-            {"role": "assistant", "content": f"⚠️ Error remoto: {e}"},
+            {
+                "role": "assistant",
+                "content": f"⚠️ Error remoto (chat): {e.__class__.__name__}: {e}",
+            },
         ]
         return updated, "", f"{e}\n{tb}"
 
@@ -223,66 +259,103 @@ def clear_chat():
     return [], "", ""
 
 # =========================
-# UI — Gradio 5
+# UI en Gradio 5
 # =========================
-with gr.Blocks(title="OpScanIA — DeepSeek-OCR + TxAgent (HF Inference)", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(
+    title="OpScanIA — DeepSeek-OCR + TxAgent (HF Inference)",
+    theme=gr.themes.Soft()
+) as demo:
     gr.Markdown(
         """
         # 📄 DeepSeek-OCR → 💬 Chat Clínico (TxAgent-T1-Llama-3.1-8B remoto)
-        1) **Sube una imagen** y corre **OCR** (imagen anotada, Markdown y texto).  
-        2) **Chatea** con **TxAgent (HF Inference)** usando automáticamente el **OCR** como contexto.  
-        *Uso educativo; no reemplaza consejo médico.*
+        1. **Sube una imagen** y corre **OCR** (imagen anotada, Markdown y texto).
+        2. **Chatea** con **TxAgent**. El chat usa automáticamente el texto detectado por OCR
+           como contexto clínico.
+        
+        ⚠ Uso educativo. No reemplaza consejo médico profesional.
         """
     )
 
+    # Estados para pasar OCR -> Chat
     ocr_md_state = gr.State("")
     ocr_txt_state = gr.State("")
 
     with gr.Row():
+        # Panel OCR
         with gr.Column(scale=1):
-            image_input = gr.Image(type="pil", label="Upload Image",
-                                   sources=["upload", "clipboard", "webcam"])
+            image_input = gr.Image(
+                type="pil",
+                label="Upload Image",
+                sources=["upload", "clipboard", "webcam"]
+            )
             model_size = gr.Dropdown(
                 choices=["Tiny", "Small", "Base", "Large", "Gundam (Recommended)"],
-                value="Gundam (Recommended)", label="Model Size"
+                value="Gundam (Recommended)",
+                label="OCR Model Size"
             )
             task_type = gr.Dropdown(
                 choices=["Free OCR", "Convert to Markdown"],
-                value="Convert to Markdown", label="Task Type"
+                value="Convert to Markdown",
+                label="OCR Task"
             )
             eval_mode_checkbox = gr.Checkbox(
                 value=True,
                 label="Evaluation mode (más rápido)",
-                info="Salida solo texto/markdown si así lo decide el backend."
+                info="Puede omitir imagen anotada y concentrarse en el texto."
             )
             submit_btn = gr.Button("Process Image", variant="primary")
 
+        # Resultados OCR
         with gr.Column(scale=2):
             with gr.Tabs():
                 with gr.TabItem("Annotated Image"):
                     output_image = gr.Image(interactive=False)
                 with gr.TabItem("Markdown Preview"):
                     output_markdown = gr.Markdown()
-                with gr.TabItem("Markdown Source / Eval Output"):
-                    output_text = gr.Textbox(lines=18, show_copy_button=True, interactive=False)
-            with gr.Row():
-                md_preview = gr.Textbox(label="Snapshot Markdown OCR", lines=8, interactive=False)
-                txt_preview = gr.Textbox(label="Snapshot Texto OCR", lines=8, interactive=False)
+                with gr.TabItem("Markdown / OCR Text"):
+                    output_text = gr.Textbox(
+                        lines=18,
+                        show_copy_button=True,
+                        interactive=False
+                    )
 
-    gr.Markdown("## Chat Clínico — TxAgent (HF Inference)")
+            with gr.Row():
+                md_preview = gr.Textbox(
+                    label="Snapshot Markdown OCR",
+                    lines=8,
+                    interactive=False
+                )
+                txt_preview = gr.Textbox(
+                    label="Snapshot Texto OCR",
+                    lines=8,
+                    interactive=False
+                )
+
+    # Panel Chat
+    gr.Markdown("## Chat Clínico — TxAgent (HF Inference / featherless-ai)")
     with gr.Row():
         with gr.Column(scale=2):
-            chatbot = gr.Chatbot(label="Asistente OCR (TxAgent remoto)", type="messages", height=420)
-            user_in = gr.Textbox(label="Mensaje",
-                                 placeholder="Escribe tu consulta… (vacío = analiza solo el OCR)",
-                                 lines=2)
+            chatbot = gr.Chatbot(
+                label="Asistente OCR (TxAgent remoto)",
+                type="messages",
+                height=420,
+            )
+            user_in = gr.Textbox(
+                label="Mensaje",
+                placeholder="Escribe tu consulta… (vacío = analiza solo el OCR)",
+                lines=2,
+            )
             with gr.Row():
                 send_btn = gr.Button("Enviar", variant="primary")
                 clear_btn = gr.Button("Limpiar")
         with gr.Column(scale=1):
-            error_box = gr.Textbox(label="Debug (si hay error)", lines=8, interactive=False)
+            error_box = gr.Textbox(
+                label="Debug (si hay error)",
+                lines=8,
+                interactive=False
+            )
 
-    # OCR
+    # Wiring OCR
     submit_btn.click(
         fn=ocr_infer,
         inputs=[image_input, model_size, task_type, eval_mode_checkbox],
@@ -293,16 +366,18 @@ with gr.Blocks(title="OpScanIA — DeepSeek-OCR + TxAgent (HF Inference)", theme
         outputs=[ocr_md_state, ocr_txt_state, md_preview, txt_preview],
     )
 
-    # Chat
+    # Wiring Chat
     send_btn.click(
         fn=chat_reply,
         inputs=[user_in, chatbot, ocr_md_state, ocr_txt_state],
-        outputs=[chatbot, user_in, error_box]
+        outputs=[chatbot, user_in, error_box],
+    )
+    clear_btn.click(
+        fn=clear_chat,
+        outputs=[chatbot, user_in, error_box],
     )
-    clear_btn.click(fn=clear_chat, outputs=[chatbot, user_in, error_box])
 
 if __name__ == "__main__":
-    # En Gradio 5 ya no existe concurrency_count en queue()
-    # Lanza directo (o usa demo.queue(max_size=…) si quisieras limitar cola).
-    # demo.queue(max_size=32)
+    # Gradio 5: sin concurrency_count en queue()
+    # demo.queue(max_size=32)  # opcional si quieres limitar cola
     demo.launch()