Update app.py

app.py

@@ -1,321 +1,57 @@
-import os
-import time
-import re
-from datetime import datetime, timezone
-from functools import lru_cache
-
-import gradio as gr
-import torch
-
-# Timezone conversion (Python 3.9+ stdlib)
-try:
-    from zoneinfo import ZoneInfo
-except Exception:
-    ZoneInfo = None  # graceful fallback to UTC
-
-# Try to import Cohere SDK if present (for hosted path)
-try:
-    import cohere  # pip install cohere
-    _HAS_COHERE = True
-except Exception:
-    _HAS_COHERE = False
-
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login, HfApi
-
-
-# -------------------
-# Configuration
-# -------------------
-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")
-)
-
-COHERE_API_KEY = os.getenv("COHERE_API_KEY")
-USE_HOSTED_COHERE = bool(COHERE_API_KEY and _HAS_COHERE)
-
-
-# -------------------
-# Helpers (used for the connection card only)
-# -------------------
-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 header(processing_time=None, user_tz: str | None = None):
-    """Only used in the connection status panel (not in chat replies)."""
-    tz_label, now_str = local_now_str(user_tz)
-    s = (
-        f"Current Date and Time ({tz_label} - YYYY-MM-DD HH:MM:SS formatted): {now_str}\n"
-        f"Current User's Login: Raj-VedAI\n"
-    )
-    if processing_time is not None:
-        s += f"Processing Time: {processing_time:.2f} seconds\n"
-    return 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)
-
-
-def is_identity_query(message: str, history) -> bool:
-    """Detects identity questions in current message or most recent user turn."""
-    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",
-    ]
-
-    def hit(text: str | None) -> bool:
-        t = (text or "").strip().lower()
-        return any(re.search(p, t) for p in patterns)
-
-    if hit(message):
-        return True
-
-    if history:
-        # Gradio history: List[Tuple[user, assistant]]
-        last_user = history[-1][0] if isinstance(history[-1], (list, tuple)) and history[-1] else None
-        if hit(last_user):
-            return True
-
-    return False
-
-
-# -------------------
-# Cohere Hosted Path
-# -------------------
-_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,
-            )
-        return _cohere_parse(resp)
-    except Exception as e:
-        return f"Error calling Cohere API: {e}"
-
-
-# -------------------
-# Local HF Path
-# -------------------
-@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."
-        )
-
-    login(token=HF_TOKEN, add_to_git_credential=False)
-
-    dtype, device_map = pick_dtype_and_map()
-    tok = AutoTokenizer.from_pretrained(
-        MODEL_ID,
-        token=HF_TOKEN,
-        use_fast=True,
-        model_max_length=4096,
-        padding_side="left",
-        trust_remote_code=True,
-    )
-    mdl = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        token=HF_TOKEN,
-        device_map=device_map,
-        low_cpu_mem_usage=True,
-        torch_dtype=dtype,
-        trust_remote_code=True,
-    )
-    if mdl.config.eos_token_id is None and tok.eos_token_id is not None:
-        mdl.config.eos_token_id = tok.eos_token_id
-    return mdl, tok
-
-
-def build_inputs(tokenizer, message, history):
-    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})
-    return tokenizer.apply_chat_template(
-        msgs, tokenize=True, add_generation_prompt=True, return_tensors="pt"
-    )
-
-
-def local_generate(model, tokenizer, input_ids, max_new_tokens=350):
-    input_ids = input_ids.to(model.device)
-    with torch.no_grad():
-        out = model.generate(
-            input_ids=input_ids,
-            max_new_tokens=max_new_tokens,
-            do_sample=True,
-            temperature=0.3,
-            top_p=0.9,
-            repetition_penalty=1.15,
-            pad_token_id=tokenizer.eos_token_id,
-            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()
-
-
-# -------------------
-# Chat callback (no header/meta in chat replies)
-# -------------------
-def chat_fn(message, history, user_tz):
-    try:
-        # Identity override → return ONLY the brand line
-        if is_identity_query(message, history):
-            return "I am ClarityOps, your strategic decision making AI partner. How can I assist you?"
-
-        if USE_HOSTED_COHERE:
-            return cohere_chat(message, history)
-
-        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}"
-
-
-# -------------------
-# Connection check (keeps header/meta)
-# -------------------
-def check_connection(user_tz=None):
-    try:
-        mode = "Cohere API (hosted)" if USE_HOSTED_COHERE else "Local HF"
-        if USE_HOSTED_COHERE:
-            return (
-                f"{header(user_tz=user_tz)}"
-                f"Connection Status: ✅ Using Cohere hosted API\n"
-                f"Mode: {mode}\n"
-                f"Model: command-r7b-12-2024\n"
-            )
-        api = HfApi(token=HF_TOKEN)
-        mi = api.model_info(MODEL_ID)
-        return (
-            f"{header(user_tz=user_tz)}"
-            f"Connection Status: ✅ Connected\n"
-            f"Mode: {mode}\n"
-            f"Model: {mi.modelId}\n"
-            f"Last Modified: {mi.lastModified}\n"
-        )
-    except Exception as e:
-        return f"{header(user_tz=user_tz)}Connection Status: ❌ Error\nDetails: {e}"
-
-
 # -------------------
 # UI
 # -------------------
 with gr.Blocks(theme=gr.themes.Default()) as demo:
-    # Hidden textbox to hold browser timezone (Gradio expects components for outputs)
+    # Hidden textbox to hold browser timezone
     tz_box = gr.Textbox(visible=False)
 
     # On load, capture browser timezone via JS and write it into tz_box
     demo.load(
-        fn=lambda tz: tz,                     # echo the JS value back to Python
-        inputs=[tz_box],                      # must provide 1 input for the lambda
-        outputs=[tz_box],                     # write result into the same hidden box
+        fn=lambda tz: tz,
+        inputs=[tz_box],
+        outputs=[tz_box],
         js="() => Intl.DateTimeFormat().resolvedOptions().timeZone"
     )
 
+    # Automatically determine connection info once tz is available
+    def model_status(user_tz):
+        try:
+            if USE_HOSTED_COHERE:
+                return "✅ Connected to: Cohere API (model: command-r7b-12-2024)"
+            api = HfApi(token=HF_TOKEN)
+            mi = api.model_info(MODEL_ID)
+            return f"✅ Connected to: Local HF model ({mi.modelId})"
+        except Exception as e:
+            return f"❌ Connection Error: {e}"
+
+    # Heading
     gr.Markdown("# Medical Decision Support AI")
 
-    with gr.Row():
-        btn = gr.Button("Check Connection Status")
-        status = gr.Textbox(label="Connection Status", lines=7, value="Click to check…")
+    # One-line status bar
+    status_line = gr.Markdown("Connecting...")
+
+    demo.load(fn=model_status, inputs=[tz_box], outputs=[status_line])
 
     gr.Markdown(
         "⚙️ First response may take a moment while the model warms up. "
-        "Currently configured to use **Cohere hosted API** if `COHERE_API_KEY` is set; "
-        "otherwise, tries **local HF**."
+       # "Currently configured to use **Cohere hosted API** if `COHERE_API_KEY` is set; "
+       # "otherwise, tries **local HF**."
     )
 
     chat = gr.ChatInterface(
         fn=chat_fn,
         type="messages",
-        additional_inputs=[tz_box],  # pass timezone into chat_fn
+        additional_inputs=[tz_box],
         description="A medical decision support system that provides healthcare-related information and decision making support.",
-        # Each example is [message, timezone]; second field is a dummy since tz comes from JS.
         examples=[
             ["What are the symptoms of hypertension?", ""],
             ["What are common drug interactions with aspirin?", ""],
             ["What are the warning signs of diabetes?", ""],
         ],
-        cache_examples=True,   # show example chips on load
+        cache_examples=True,
     )
 
-    # Wire timezone into the connection check as well
-    btn.click(fn=check_connection, inputs=tz_box, outputs=status)
-
 if __name__ == "__main__":
     demo.launch()
 
+