Update app.py

app.py

@@ -6,15 +6,15 @@ from functools import lru_cache
 import gradio as gr
 import torch
 
-# Timezone conversion (Python 3.9+ stdlib)
+# Timezone (Python 3.9+)
 try:
     from zoneinfo import ZoneInfo
 except Exception:
-    ZoneInfo = None  # graceful fallback to UTC
+    ZoneInfo = None
 
-# Try Cohere SDK if present (for hosted path)
+# Cohere SDK
 try:
-    import cohere  # pip install cohere
+    import cohere
     _HAS_COHERE = True
 except Exception:
     _HAS_COHERE = False
@@ -23,124 +23,84 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login, HfApi
 
 # -------------------
-# Configuration
+# Config
 # -------------------
 MODEL_ID = os.getenv("MODEL_ID", "CohereLabs/c4ai-command-r7b-12-2024")
-
-HF_TOKEN = (
-    os.getenv("HUGGINGFACE_HUB_TOKEN")  # official Spaces name
-    or os.getenv("HF_TOKEN")
-)
-
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN") or os.getenv("HF_TOKEN")
 COHERE_API_KEY = os.getenv("COHERE_API_KEY")
 USE_HOSTED_COHERE = bool(COHERE_API_KEY and _HAS_COHERE)
 
 # -------------------
-# Helpers (status only)
+# Helpers
 # -------------------
-def local_now_str(user_tz: str | None) -> tuple[str, str]:
-    """Returns (label, formatted_time). Falls back to UTC if tz missing/invalid."""
-    label = "UTC"
-    dt = datetime.now(timezone.utc)
-    if user_tz and ZoneInfo is not None:
-        try:
-            tz = ZoneInfo(user_tz)
-            dt = datetime.now(tz)
-            label = user_tz
-        except Exception:
-            dt = datetime.now(timezone.utc)
-            label = "UTC"
-    return label, dt.strftime("%Y-%m-%d %H:%M:%S")
-
-
 def pick_dtype_and_map():
     if torch.cuda.is_available():
         return torch.float16, "auto"
     if torch.backends.mps.is_available():
         return torch.float16, {"": "mps"}
-    return torch.float32, "cpu"  # CPU path (likely too big for R7B)
+    return torch.float32, "cpu"
 
-def is_identity_query(message: str, history) -> bool:
-    """Detects identity questions in current message or most recent user turn."""
+def is_identity_query(message, history):
     patterns = [
-        r"\bwho\s+are\s+you\b",
-        r"\bwhat\s+are\s+you\b",
-        r"\bwhat\s+is\s+your\s+name\b",
-        r"\bwho\s+is\s+this\b",
-        r"\bidentify\s+yourself\b",
-        r"\btell\s+me\s+about\s+yourself\b",
-        r"\bdescribe\s+yourself\b",
-        r"\band\s+you\s*\?\b",
-        r"\byour\s+name\b",
-        r"\bwho\s+am\s+i\s+chatting\s+with\b",
+        r"\bwho\s+are\s+you\b", r"\bwhat\s+are\s+you\b",
+        r"\bwhat\s+is\s+your\s+name\b", r"\bwho\s+is\s+this\b",
+        r"\bidentify\s+yourself\b", r"\btell\s+me\s+about\s+yourself\b",
+        r"\bdescribe\s+yourself\b", r"\band\s+you\s*\?\b",
+        r"\byour\s+name\b", r"\bwho\s+am\s+i\s+chatting\s+with\b"
     ]
-    def hit(text: str | None) -> bool:
-        t = (text or "").strip().lower()
-        return any(re.search(p, t) for p in patterns)
-    if hit(message):
+    def match(t):
+        return any(re.search(p, (t or "").strip().lower()) for p in patterns)
+    if match(message):
         return True
     if history:
-        last_user = history[-1][0] if isinstance(history[-1], (list, tuple)) and history[-1] else None
-        if hit(last_user):
+        last_user = history[-1][0] if isinstance(history[-1], (list, tuple)) else None
+        if match(last_user):
             return True
     return False
 
 # -------------------
-# Cohere Hosted Path
+# Cohere Hosted
 # -------------------
 _co_client = None
 if USE_HOSTED_COHERE:
     _co_client = cohere.Client(api_key=COHERE_API_KEY)
 
 def _cohere_parse(resp):
-    # v5+ responses.create
     if hasattr(resp, "output_text") and resp.output_text:
         return resp.output_text.strip()
     if getattr(resp, "message", None) and getattr(resp.message, "content", None):
         for p in resp.message.content:
             if hasattr(p, "text") and p.text:
                 return p.text.strip()
-    # v4 chat
     if hasattr(resp, "text") and resp.text:
         return resp.text.strip()
     return "Sorry, I couldn't parse the response from Cohere."
 
 def cohere_chat(message, history):
     try:
-        # Prefer modern API
-        try:
-            msgs = []
-            for u, a in (history or []):
-                msgs.append({"role": "user", "content": u})
-                msgs.append({"role": "assistant", "content": a})
-            msgs.append({"role": "user", "content": message})
-            resp = _co_client.responses.create(
-                model="command-r7b-12-2024",
-                messages=msgs,
-                temperature=0.3,
-                max_tokens=350,
-            )
-        except Exception:
-            # Fallback to older chat API
-            resp = _co_client.chat(
-                model="command-r7b-12-2024",
-                message=message,
-                temperature=0.3,
-                max_tokens=350,
-            )
+        msgs = []
+        for u, a in (history or []):
+            msgs.append({"role": "user", "content": u})
+            msgs.append({"role": "assistant", "content": a})
+        msgs.append({"role": "user", "content": message})
+        resp = _co_client.responses.create(
+            model="command-r7b-12-2024",
+            messages=msgs,
+            temperature=0.3,
+            max_tokens=350,
+        )
         return _cohere_parse(resp)
     except Exception as e:
         return f"Error calling Cohere API: {e}"
 
 # -------------------
-# Local HF Path
+# Local HF Model
 # -------------------
 @lru_cache(maxsize=1)
 def load_local_model():
     if not HF_TOKEN:
         raise RuntimeError(
-            "HUGGINGFACE_HUB_TOKEN (or HF_TOKEN) is not set. "
-            "Either set it, or provide COHERE_API_KEY to use Cohere's hosted API."
+            "HUGGINGFACE_HUB_TOKEN is not set."
         )
     login(token=HF_TOKEN, add_to_git_credential=False)
     dtype, device_map = pick_dtype_and_map()
@@ -188,11 +148,10 @@ def local_generate(model, tokenizer, input_ids, max_new_tokens=350):
             eos_token_id=tokenizer.eos_token_id,
         )
     gen_only = out[0, input_ids.shape[-1]:]
-    text = tokenizer.decode(gen_only, skip_special_tokens=True)
-    return text.strip()
+    return tokenizer.decode(gen_only, skip_special_tokens=True).strip()
 
 # -------------------
-# Chat callback (no meta in replies)
+# Chat Function
 # -------------------
 def chat_fn(message, history, user_tz):
     try:
@@ -203,97 +162,39 @@ def chat_fn(message, history, user_tz):
         model, tokenizer = load_local_model()
         inputs = build_inputs(tokenizer, message, history)
         return local_generate(model, tokenizer, inputs, max_new_tokens=350)
-    except RuntimeError as e:
-        emsg = str(e)
-        if "out of memory" in emsg.lower() or "cuda" in emsg.lower():
-            return "Local load likely OOM. Use a GPU Space or set COHERE_API_KEY to run via Cohere hosted API."
-        return f"Error during chat: {e}"
     except Exception as e:
-        return f"Error during chat: {e}"
+        return f"Error: {e}"
 
 # -------------------
-# Theme & Styles (compatible with broad Gradio versions)
+# Theme & CSS
 # -------------------
 theme = gr.themes.Soft(
     primary_hue="teal",
     neutral_hue="slate",
     radius_size=gr.themes.sizes.radius_lg,
-).set(
-    shadow_drop="0 6px 24px rgba(0,0,0,.06)",
-    shadow_spread="0 2px 8px rgba(0,0,0,.04)",
 )
 
 custom_css = """
 :root {
-  --brand-bg: #e6f7f8;        /* soft medical teal */
-  --brand-card: #ffffff;
-  --brand-text: #0f172a;      /* slate-900 */
-  --brand-subtle: #475569;    /* slate-600 */
-  --brand-accent: #0d9488;    /* teal-600 */
-  --brand-border: #cbd5e1;    /* slate-300 */
+  --brand-bg: #e6f7f8; /* soft medical teal */
+  --brand-accent: #0d9488; /* teal-600 */
+  --brand-text-light: #ffffff;
 }
 
-/* Page background */
 .gradio-container {
   background: var(--brand-bg);
-  color: var(--brand-text);
 }
 
-/* Title */
-h1, .prose h1 {
-  color: var(--brand-text);
+h1 {
+  color: #0f172a;
   font-weight: 700;
-  letter-spacing: -0.01em;
-  margin-bottom: 0.25rem !important;
-  font-size: 28px !important; /* set via CSS for compatibility */
-}
-
-/* Chat bubbles */
-.message.user {
-  background: var(--brand-accent) !important; /* teal bubble */
-  color: #ffffff !important;                  /* white text */
-}
-.message.bot {
-  background: var(--brand-card) !important;   /* white bubble */
-  color: var(--brand-text) !important;        /* dark text */
-}
-
-/* Status badge wrapper */
-.status-wrap {
-  display: flex;
-  align-items: center;
-  gap: .5rem;
-  margin-bottom: 0.75rem;
+  font-size: 28px !important;
 }
 
-/* Badge */
-.badge {
-  display: inline-flex;
-  align-items: center;
-  gap: .5rem;
-  padding: .45rem .75rem;
-  border-radius: 999px;
-  border: 1px solid var(--brand-border);
-  background: #ecfdf5; /* green-50 */
-  color: #065f46;      /* green-800 */
-  font-weight: 600;
-  font-size: 14px;
-}
-
-/* Helper text */
-.helper {
-  color: var(--brand-subtle);
-  margin: .25rem 0 1rem 0;
-}
-
-/* Card rounding */
-.block, .gr-box, .gr-panel, .gr-group, .gr-form, .gradio-container .form {
-  border-radius: 16px !important;
-}
-
-/* Inputs */
-textarea, input, .gr-input {
-  border-radius: 12px !important;
+/* Both bot and user bubbles teal with white text */
+.message.user, .message.bot {
+  background: var(--brand-accent) !important;
+  color: var(--brand-text-light) !important;
 }
 """
 
@@ -301,57 +202,16 @@ textarea, input, .gr-input {
 # UI
 # -------------------
 with gr.Blocks(theme=theme, css=custom_css) as demo:
-    # Hidden textbox to hold browser timezone
     tz_box = gr.Textbox(visible=False)
+    demo.load(lambda tz: tz, inputs=[tz_box], outputs=[tz_box],
+              js="() => Intl.DateTimeFormat().resolvedOptions().timeZone")
 
-    # Capture browser timezone via JS and store in tz_box
-    demo.load(
-        fn=lambda tz: tz,   # echo JS value
-        inputs=[tz_box],
-        outputs=[tz_box],
-        js="() => Intl.DateTimeFormat().resolvedOptions().timeZone"
-    )
-
-    # Model status (auto, one-line badge)
-    def model_status(_user_tz):
-        try:
-            if USE_HOSTED_COHERE:
-                return (
-                    '<div class="status-wrap">'
-                    '<span class="badge">✅ Connected • Cohere API — model: '
-                    '<strong>command-r7b-12-2024</strong></span></div>'
-                )
-            api = HfApi(token=HF_TOKEN)
-            mi = api.model_info(MODEL_ID)
-            return (
-                '<div class="status-wrap">'
-                f'<span class="badge">✅ Connected • Local HF — model: '
-                f'<strong>{mi.modelId}</strong></span></div>'
-            )
-        except Exception as e:
-            return (
-                '<div class="status-wrap">'
-                f'<span class="badge" style="background:#fff7ed;color:#9a3412;border-color:#fed7aa;">'
-                f'⚠️ Connection Issue — {str(e)}</span></div>'
-            )
-
-    # Header + status
     gr.Markdown("# Medical Decision Support AI")
-    status_line = gr.HTML("<div class='status-wrap'><span class='badge'>Connecting…</span></div>")
-    demo.load(fn=model_status, inputs=[tz_box], outputs=[status_line])
-
-    # Helper text
-    gr.Markdown(
-        "<div class='helper'>Designed for healthcare executives: concise, reliable decision support. "
-        "First response may take a moment while the model warms up.</div>"
-    )
 
-    # Chat
     gr.ChatInterface(
         fn=chat_fn,
         type="messages",
-        additional_inputs=[tz_box],  # pass timezone into chat_fn (future use)
-        description="",
+        additional_inputs=[tz_box],
         examples=[
             ["What are the symptoms of hypertension?", ""],
             ["What are common drug interactions with aspirin?", ""],