Upload app.py with huggingface_hub

app.py

@@ -9,15 +9,89 @@ import json
 import os
 import tempfile
 import threading
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from pathlib import Path
-from backend import translate_to_english, call_model_a, call_model_b
+from backend import ModelResult, translate_to_english, call_model_a, call_model_b
 
 FEEDBACK_FILE = Path(__file__).parent / "feedback_data.json"
 HF_DATASET_REPO = "Vrda/im-error-check-data"
 HF_DATASET_FILE = "feedback_data.json"
 
-_DEEPSEEK_RESULTS: dict[str, dict] = {}
+
+@st.cache_resource
+def get_deepseek_job_manager():
+    return {
+        "executor": ThreadPoolExecutor(max_workers=2),
+        "jobs": {},
+        "lock": threading.Lock(),
+    }
+
+
+def cleanup_deepseek_jobs(max_age_seconds: int = 1800):
+    manager = get_deepseek_job_manager()
+    now = time.time()
+    stale_job_ids = []
+    with manager["lock"]:
+        for job_id, job in manager["jobs"].items():
+            if now - job["created_at"] > max_age_seconds:
+                stale_job_ids.append(job_id)
+        for job_id in stale_job_ids:
+            manager["jobs"].pop(job_id, None)
+
+
+def submit_deepseek_job(job_id: str, english_text: str):
+    manager = get_deepseek_job_manager()
+    future = manager["executor"].submit(call_model_a, english_text)
+    with manager["lock"]:
+        manager["jobs"][job_id] = {
+            "future": future,
+            "created_at": time.time(),
+        }
+
+
+def get_deepseek_job_info(job_id: str):
+    if not job_id:
+        return None
+    manager = get_deepseek_job_manager()
+    with manager["lock"]:
+        job = manager["jobs"].get(job_id)
+        if not job:
+            return None
+        return {
+            "created_at": job["created_at"],
+            "done": job["future"].done(),
+        }
+
+
+def consume_deepseek_job_result(job_id: str) -> ModelResult | None:
+    if not job_id:
+        return None
+    manager = get_deepseek_job_manager()
+    with manager["lock"]:
+        job = manager["jobs"].get(job_id)
+    if not job:
+        return None
+
+    future = job["future"]
+    if not future.done():
+        return None
+
+    try:
+        result = future.result()
+    except Exception as exc:
+        result = ModelResult(
+            model_name="DeepSeek Reasoner",
+            raw_response="",
+            success=False,
+            error_message=f"Background job failed: {exc}",
+            latency_seconds=0.0,
+        )
+
+    with manager["lock"]:
+        manager["jobs"].pop(job_id, None)
+
+    return result
 
 # -------------------------------------------------------------------------
 # Feedback persistence (local + HF Hub sync)
@@ -186,12 +260,39 @@ for key, default in [
     ("model_b_result", None),
     ("translation_latency", 0),
     ("total_elapsed", 0),
+    ("analysis_started_at", 0.0),
+    ("deepseek_job_id", None),
     ("run_analysis", False),
     ("physician_id", ""),
 ]:
     if key not in st.session_state:
         st.session_state[key] = default
 
+if "session_key" not in st.session_state:
+    import uuid
+
+    st.session_state.session_key = str(uuid.uuid4())
+
+cleanup_deepseek_jobs()
+
+
+@st.fragment(run_every=5)
+def poll_deepseek_job():
+    job_id = st.session_state.deepseek_job_id
+    if not job_id or st.session_state.model_a_result is not None:
+        return
+
+    result = consume_deepseek_job_result(job_id)
+    if result is None:
+        return
+
+    st.session_state.model_a_result = result
+    st.session_state.deepseek_job_id = None
+    st.session_state.total_elapsed = round(
+        time.time() - st.session_state.analysis_started_at, 2
+    )
+    st.rerun()
+
 
 def load_sample():
     st.session_state.input_text = SAMPLE
@@ -252,24 +353,13 @@ st.button("Analyze", type="primary", on_click=trigger_analysis)
 # Run analysis (progressive: show GPT-OSS first, DeepSeek when ready)
 # -------------------------------------------------------------------------
 
-def _run_deepseek_background(session_key: str, english_text: str):
-    """Background thread: calls DeepSeek and stores result in module-level dict."""
-    result = call_model_a(english_text)
-    _DEEPSEEK_RESULTS[session_key] = result
-
-if "session_key" not in st.session_state:
-    import uuid
-    st.session_state.session_key = str(uuid.uuid4())
-
 if st.session_state.run_analysis and st.session_state.input_text.strip():
     st.session_state.run_analysis = False
     st.session_state.model_a_result = None
     st.session_state.model_b_result = None
     st.session_state.total_elapsed = 0
-    st.session_state._analysis_start = time.time()
-
-    skey = st.session_state.session_key
-    _DEEPSEEK_RESULTS.pop(skey, None)
+    st.session_state.analysis_started_at = time.time()
+    st.session_state.deepseek_job_id = None
 
     with st.spinner("Translating discharge letter..."):
         t0 = time.time()
@@ -278,16 +368,17 @@ if st.session_state.run_analysis and st.session_state.input_text.strip():
 
     english = st.session_state.translated_text
 
-    thread = threading.Thread(
-        target=_run_deepseek_background, args=(skey, english), daemon=True
-    )
-    thread.start()
+    job_id = f"{st.session_state.session_key}:{int(time.time() * 1000)}"
+    submit_deepseek_job(job_id, english)
+    st.session_state.deepseek_job_id = job_id
 
     with st.spinner("GPT-OSS-120B responding (~5s)..."):
         st.session_state.model_b_result = call_model_b(english)
 
     st.rerun()
 
+poll_deepseek_job()
+
 # -------------------------------------------------------------------------
 # Helper: render a model's output
 # -------------------------------------------------------------------------
@@ -405,16 +496,13 @@ if has_any_result:
         if st.session_state.model_a_result is not None:
             render_model_output(st.session_state.model_a_result, "model-header-a")
         else:
-            @st.fragment(run_every=5)
-            def _poll_deepseek():
-                skey = st.session_state.session_key
-                if skey in _DEEPSEEK_RESULTS:
-                    st.session_state.model_a_result = _DEEPSEEK_RESULTS.pop(skey)
-                    st.session_state.total_elapsed = round(
-                        time.time() - st.session_state._analysis_start, 2
-                    )
-                    st.rerun()
-                elapsed = round(time.time() - st.session_state._analysis_start)
+            job_info = get_deepseek_job_info(st.session_state.deepseek_job_id)
+            if job_info is None:
+                st.warning(
+                    "DeepSeek job is no longer active. Click `Analyze` to run it again."
+                )
+            else:
+                elapsed = round(time.time() - job_info["created_at"])
                 st.markdown(
                     '<div style="background:#1e293b; border:2px dashed #475569; '
                     'border-radius:8px; padding:2rem; text-align:center; color:#e2e8f0;">'
@@ -424,7 +512,7 @@ if has_any_result:
                     "</div>",
                     unsafe_allow_html=True,
                 )
-            _poll_deepseek()
+                st.caption("Checking DeepSeek status every 5 seconds.")
 
     # -----------------------------------------------------------------
     # Feedback