Update app.py

app.py

@@ -1,141 +1,164 @@
 import gradio as gr
-from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
+import torch
 import json
 import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+from huggingface_hub import InferenceClient
 
 # ==============================================================================
 # 1. CONFIGURATION
 # ==============================================================================
-HF_TOKEN = os.getenv("HF_TOKEN")
-PROJECT_TITLE = "The Janus Interface: Local Air-Gap Architecture"
+# NOTE: You must set 'HF_TOKEN' in your Hugging Face Space Secrets!
+HF_TOKEN = os.getenv("HF_TOKEN") 
 
-# GGUF Model (The Unified Adapter you converted)
-REPO_ID = "st192011/janus-unified-gguf"
-FILENAME = "janus_unified_q4_k_m.gguf" 
+PROJECT_TITLE = "The Janus Interface: Semantic Decoupling Architecture"
+
+# Models
+# We use the official Microsoft repo for CPU compatibility
+BASE_MODEL_ID = "microsoft/Phi-3.5-mini-instruct" 
+ADAPTER_ID = "st192011/janus-gold-lora" 
+CLOUD_MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct" 
 
 # ==============================================================================
-# 2. ENGINE INITIALIZATION
+# 2. ENGINE INITIALIZATION (CPU Optimized)
 # ==============================================================================
-print("⏳ Downloading GGUF Model...")
-try:
-    model_path = hf_hub_download(
-        repo_id=REPO_ID, 
-        filename=FILENAME, 
-        token=HF_TOKEN
-    )
-    print(f"✅ Model downloaded to: {model_path}")
-except Exception as e:
-    print(f"❌ Error downloading model: {e}")
-    raise e
+print("⏳ Initializing Neural Backbone (CPU Mode)...")
 
-print("⏳ Loading Llama.cpp Engine...")
 try:
-    # n_ctx=4096 matches training context
-    llm = Llama(
-        model_path=model_path,
-        n_ctx=4096,
-        n_threads=2, # Optimized for Free Tier
-        verbose=False
+    # Load Tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+    
+    # Load Base Model (bfloat16 saves RAM on Free Tier Spaces)
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL_ID,
+        torch_dtype=torch.bfloat16, 
+        device_map="cpu",
+        trust_remote_code=True
     )
+    
+    # Load Adapter
+    print(f"⏳ Mounting Janus Adapter ({ADAPTER_ID})...")
+    model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
+    model.eval() # Set to inference mode
     print("✅ System Online.")
+
 except Exception as e:
-    print(f"❌ Error initializing Llama.cpp: {e}")
+    print(f"❌ Error loading model: {e}")
     raise e
 
+# Cloud Client
+hf_client = InferenceClient(model=CLOUD_MODEL_ID, token=HF_TOKEN)
+
 # ==============================================================================
 # 3. KERNEL LOGIC
 # ==============================================================================
 
 def clean_output(text):
     """Sanitizes output to prevent chain-reaction failures."""
+    # Remove special tokens
     clean = text.replace("<|end|>", "").replace("<|endoftext|>", "")
+    
+    # Remove conversational filler lines
     if "Output:" in clean: clean = clean.split("Output:")[-1]
     
     lines = clean.split('\n')
+    # Keep lines that look like protocol code or normal text, remove "Here is..."
     valid_lines = [line for line in lines if "Note" not in line and "Here is" not in line]
     return " ".join(valid_lines).strip()
 
-def generate_response(system_prompt, user_input, max_tokens=512):
-    """Helper for GGUF Generation"""
-    prompt = f"<|system|>\n{system_prompt}<|end|>\n<|user|>\n{user_input}<|end|>\n<|assistant|>\n"
-    
-    output = llm(
-        prompt,
-        max_tokens=max_tokens,
-        temperature=0.1,
-        top_p=0.95,
-        stop=["<|end|>", "<|endoftext|>"],
-        echo=False
-    )
-    return output['choices'][0]['text']
-
 def kernel_scout(raw_input):
-    """Role: The Scout (Sanitizer)"""
+    """Mode A: Local Logic Extraction"""
     try:
-        sys_prompt = """SYSTEM_ROLE: Janus Extractor.
+        prompt = f"""<|system|>
+SYSTEM_ROLE: Janus Extractor.
 TASK: Refactor clinical notes into JanusScript Logic.
 SYNTAX: Object.action(params) -> Result.
 OBJECTS: Hx, Sx, Dx, Tx, Lab, Crs, Plan.
-CONSTRAINTS: No PII. Use relative time (Day1, Day2)."""
+CONSTRAINTS: No PII. Use relative time (Day1, Day2).
+<|end|>
+<|user|>
+RAW NOTE:
+{raw_input}<|end|>
+<|assistant|>"""
+        
+        inputs = tokenizer(prompt, return_tensors="pt")
         
-        user_input = f"RAW NOTE:\n{raw_input}"
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs, 
+                max_new_tokens=256, 
+                temperature=0.1,
+                do_sample=True
+            )
         
-        raw_output = generate_response(sys_prompt, user_input, max_tokens=256)
+        text = tokenizer.batch_decode(outputs)[0]
+        raw_output = text.split("<|assistant|>")[-1]
         return clean_output(raw_output)
-    except Exception as e: return f"Scout Error: {str(e)}"
+    except Exception as e: return f"Error: {str(e)}"
+
+def kernel_cloud_expert(scenario_prompt):
+    """Mode B: Cloud Bridge"""
+    try:
+        sys_prompt = """You are a Clinical Logic Engine.
+Task: Convert the scenario into 'JanusScript' code.
+Syntax: Object.action(parameter);
+Objects: Dx, Sx, Tx, Lab, Plan.
+Rules: No PII. Use PascalCase.
+
+Example:
+Input: Pt has pneumonia. Given antibiotics.
+Output: Dx(Pneumonia); Sx(Fever+Cough); Tx(Meds).action(Antibiotics); Plan(Discharge.Home);"""
+        
+        messages = [
+            {"role": "system", "content": sys_prompt},
+            {"role": "user", "content": f"Input: {scenario_prompt}"}
+        ]
+        
+        response = hf_client.chat_completion(messages, max_tokens=512, temperature=0.1)
+        return clean_output(response.choices[0].message.content)
+    except Exception as e: return f"API Error: {str(e)}"
 
 def kernel_vault(protocol, secure_json):
-    """Role: The Vault (Reconstructor)"""
+    """Shared Terminal: Reconstruction"""
     try:
         try: db_str = json.dumps(json.loads(secure_json), ensure_ascii=False)
         except: return "❌ Error: Invalid JSON."
 
-        sys_prompt = """SYSTEM_ROLE: Janus Constructor.
+        prompt = f"""<|system|>
+SYSTEM_ROLE: Janus Constructor.
 TASK: Interpret JanusScript and PrivateDB to write Discharge Summary.
-TEMPLATE: Header -> Dates -> History -> Hospital Course -> Plan."""
+TEMPLATE: Header -> Dates -> History -> Hospital Course -> Plan.
+<|end|>
+<|user|>
+PROTOCOL:
+{protocol}
+
+PRIVATE_DB:
+{db_str}<|end|>
+<|assistant|>"""
+
+        inputs = tokenizer(prompt, return_tensors="pt")
         
-        user_input = f"PROTOCOL:\n{protocol}\n\nPRIVATE_DB:\n{db_str}"
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs, 
+                max_new_tokens=1024, 
+                temperature=0.1, 
+                repetition_penalty=1.05,
+                do_sample=True
+            )
         
-        raw_output = generate_response(sys_prompt, user_input, max_tokens=1024)
-        return clean_output(raw_output)
-    except Exception as e: return f"Vault Error: {str(e)}"
+        text = tokenizer.batch_decode(outputs)[0]
+        doc = text.split("<|assistant|>")[-1].replace("<|end|>", "").replace("<|endoftext|>", "").strip()
+        return doc
+    except Exception as e: return f"Error: {str(e)}"
 
 # ==============================================================================
-# 4. TECHNICAL REPORT (Corrected)
-# ==============================================================================
-report_md = """
-# 🏛️ The Janus Interface: Technical Analysis
-**Research Prototype v2.0**
-
-### 1. Architectural Goal: Zero-Trust Processing
-The objective is to process sensitive clinical data using Generative AI without exposing **Private Identity Information (PII)** to the reasoning engine. We achieve this via **Semantic Decoupling**:
-1.  **The Scout (Logic Layer):** Extracts clinical entities (`DX`, `PROC`) into an anonymous Protocol.
-2.  **The Air Gap:** Only the Protocol string crosses the boundary.
-3.  **The Vault (Context Layer):** Reconstructs the document by merging the Protocol with a local SQL Database.
-
-### 2. Training Methodology
-*   **Model:** Microsoft Phi-3.5-mini (3.8B Parameters).
-*   **Technique:** Multi-Task Instruction Tuning using Unsloth (QLoRA + DoRA).
-*   **Dataset:** 306 Fully Synthesized Clinical Scenarios.
-*   **Strategy:** The model was trained to switch behaviors based on the `SYSTEM_ROLE` prompt ("Extractor" vs "Constructor").
-
-### 3. Data Engineering: Synthetic Alignment Pipeline
-To eliminate the hallucinations observed in earlier iterations (where the model confused patient gender or age), we utilized a **Fully Synthetic Data Pipeline**:
-1.  **Generation:** We used Large Language Models (LLMs) to generate 306 unique clinical scenarios.
-2.  **Alignment:** Unlike real-world datasets where headers might be missing or mismatched, this synthetic dataset ensured strict mathematical consistency between the Identity Header (Age/Sex) and the Clinical Narrative.
-3.  **Result:** This alignment forced the model to respect the Database variables during inference, successfully resolving the "Gender Mismatch" and "Age Hallucination" issues.
-
-### 4. Limitations & Constraints
-*   **Vocabulary:** The model is trained on a specific subset of medical cases (Surgery, Internal Medicine). It may struggle with out-of-distribution specialties (e.g., Oncology, Psychiatry) not present in the 306 samples.
-*   **Compression Loss:** The "Scout" layer intentionally compresses text. Nuanced emotional context or very specific non-medical details in the raw note may be lost during the conversion to JanusScript.
-"""
-
-# ==============================================================================
-# 5. DEMO SAMPLES
+# 4. DEMO SAMPLES
 # ==============================================================================
 
-# Case: Appendicitis (Matches Training Distribution)
+# Case 1: Appendicitis (Local)
 sample_note = """Pt ID 8899-A.
 History: 28yo male presented with RLQ pain & fever. 
 Workup: CT scan confirmed acute appy.
@@ -153,43 +176,133 @@ sample_db = """{
     "prov_specialty": "General Surgery"
 }"""
 
+# Case 2: Sepsis (Cloud)
+sample_scenario = """Patient admitted for Urosepsis.
+Culture: E. coli resistant to Cipro.
+Treatment: Started on Zosyn IV. Transferred to ICU for one day for hypotension.
+Transition: Switched to oral Augmentin.
+Outcome: Stable, Afebrile. Discharge to finish 14 day course."""
+
+sample_db_cloud = """{
+    "pt_name": "Sarah Connor",
+    "pt_mrn": "SKY-NET",
+    "pt_dob": "1965-05-10",
+    "pt_sex": "F",
+    "adm_date": "2025-12-01",
+    "dis_date": "2025-12-05",
+    "prov_attending": "Dr. Silberman",
+    "prov_specialty": "Internal Medicine"
+}"""
+
 # ==============================================================================
-# 6. GRADIO LAUNCHER
+# 5. TECHNICAL REPORT
 # ==============================================================================
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="zinc"), title=PROJECT_TITLE) as demo:
+report_md = """
+# 🏛️ The Janus Interface: Research & Technical Analysis
+**Project Status:** Research Prototype v2.0 (Gold Standard)
+
+---
+
+### 1. Research Motivation: The Privacy-Utility Paradox
+In regulated domains (Healthcare, Legal, Finance), Generative AI adoption is stalled by a fundamental conflict:
+*   **Utility:** Large Cloud Models (GPT-4, Claude) offer superior reasoning but require sending data off-premise.
+*   **Privacy:** Local Small Models (SLMs) ensure data sovereignty but often lack deep domain knowledge.
+*   **The Solution:** **Semantic Decoupling**. We propose separating the **"Logic"** of a case from the **"Identity"** of the subject.
+
+### 2. Architectural Design: The Twin-Protocol
+The system utilizes a **Multi-Task Adapter** trained to switch between two distinct cognitive modes based on the System Prompt.
+
+#### **Mode A: The Scout (Logic Extractor)**
+*   **Function:** Reads raw, messy clinical notes.
+*   **Constraint:** Trained via Loss Masking to extract *only* clinical entities (`Dx`, `Tx`, `Plan`) into a sanitized code string called **JanusScript**.
+*   **Security:** It treats names, dates, and locations as noise to be discarded.
+
+#### **Mode B: The Cloud Bridge (Knowledge Injection)**
+*   **Function:** Allows an external Cloud LLM to reason about a generic, anonymized scenario.
+*   **Innovation:** The Cloud Model generates the **JanusScript** code. This code acts as a firewall—no PII ever leaves the local environment, but the *intelligence* of the cloud is captured in the script.
+
+#### **The Vault (Reconstructor)**
+*   **Function:** A secure, offline engine that accepts the JanusScript and a Local SQL Database record.
+*   **Output:** It merges the abstract logic with the concrete identity to generate the final, human-readable document.
+
+---
+
+### 3. Data Engineering: The "Gold Standard" Pipeline
+To achieve high fidelity without using private patient data, we developed a **Synthesized Data Pipeline**:
+
+1.  **Synthesis:** We generated **306 high-quality clinical scenarios** using Large Language Models (LLMs).
+2.  **Alignment:** Unlike previous iterations where headers were random, this dataset ensured strict mathematical alignment between the Identity Header (Age/Sex) and the Clinical Narrative.
+3.  **Result:** This eliminated the "hallucination" issues seen in earlier tests where the model would confuse patient gender or age due to conflicting training signals.
+
+### 4. Training Methodology
+*   **Base Model:** Microsoft Phi-3.5-mini-instruct (3.8B Parameters).
+*   **Framework:** **Unsloth** (Optimized QLoRA).
+*   **Technique:** **DoRA (Weight-Decomposed Low-Rank Adaptation)**.
+    *   *Why DoRA?* Standard LoRA struggles with strict syntax/coding tasks. DoRA updates both magnitude and direction vectors, allowing the model to learn the strict `JanusScript` grammar effectively.
+*   **Loss Masking:** We used `train_on_responses_only`. The model was **never** trained on the input text, only on the output. This prevents the model from memorizing patient PII from the training set.
+*   **Hyperparameters:** Rank 16, Alpha 16, Learning Rate 2e-4, **2 Epochs** (approx 78 steps used for final checkpoint).
+
+### 5. Results & Conclusion
+*   **Zero-Trust Validation:** The "Vault" successfully reconstructs documents using *only* the database for identity.
+*   **Semantic Expansion:** The model demonstrates the ability to take a concise code (`Dx(Pneumonia)`) and expand it into fluent medical narrative ("Patient presented with symptoms consistent with Pneumonia...").
+"""
+
+# ==============================================================================
+# 6. LAUNCHER
+# ==============================================================================
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald"), title=PROJECT_TITLE) as demo:
     gr.Markdown(f"# 🏛️ {PROJECT_TITLE}")
     
     with gr.Tabs():
         
-        # --- TAB 1: DEMO ---
-        with gr.TabItem("🛡️ Live Demo"):
-            gr.Markdown("### Zero-Trust Clinical Documentation")
-            
+        # --- TAB 1 ---
+        with gr.TabItem("🛡️ Mode A: Local Air-Gap"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    inp_scout = gr.Textbox(label="1. Raw Sensitive Note (Unsecured)", lines=12, value=sample_note)
-                    btn_scout = gr.Button("Execute Scout (Extract Logic) ➔", variant="primary")
+                    inp_a = gr.Textbox(label="Raw Sensitive Note", lines=12, value=sample_note)
+                    btn_a = gr.Button("Execute Scout (Local) ➔", variant="primary")
                 
                 with gr.Column(scale=1):
-                    out_proto = gr.Textbox(label="2. Intermediate Protocol (The Air-Gap)", lines=6, interactive=True,
-                                         info="This code string contains NO identity data.")
+                    out_proto_a = gr.Textbox(label="JanusScript Protocol", lines=6, interactive=True)
             
             gr.Markdown("---")
             
             with gr.Row():
                 with gr.Column(scale=1):
-                    inp_db = gr.Textbox(label="3. Secure Identity Record (Local DB)", lines=12, value=sample_db)
-                    btn_vault = gr.Button("Execute Vault (Reconstruct) ➔", variant="secondary")
+                    inp_db_a = gr.Textbox(label="Secure Identity Record", lines=12, value=sample_db)
+                    btn_final_a = gr.Button("Execute Vault (Local) ➔", variant="secondary")
+                
+                with gr.Column(scale=1):
+                    out_final_a = gr.Textbox(label="Output: Reconstructed Document", lines=25)
+            
+            btn_a.click(kernel_scout, inputs=inp_a, outputs=out_proto_a)
+            btn_final_a.click(kernel_vault, inputs=[out_proto_a, inp_db_a], outputs=out_final_a)
+
+        # --- TAB 2 ---
+        with gr.TabItem("🧠 Mode B: Cloud Bridge"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    inp_b = gr.Textbox(label="Clinical Scenario (Anonymized)", lines=12, value=sample_scenario)
+                    btn_b = gr.Button("Execute Cloud API (Llama-3) ➔", variant="primary")
                 
                 with gr.Column(scale=1):
-                    out_final = gr.Textbox(label="4. Final Document", lines=25)
+                    out_proto_b = gr.Textbox(label="JanusScript Protocol", lines=6, interactive=True)
             
-            # Wiring
-            btn_scout.click(kernel_scout, inputs=inp_scout, outputs=out_proto)
-            btn_vault.click(kernel_vault, inputs=[out_proto, inp_db], outputs=out_final)
+            gr.Markdown("---")
+            
+            with gr.Row():
+                with gr.Column(scale=1):
+                    inp_db_b = gr.Textbox(label="Secure Identity Record", lines=12, value=sample_db_cloud)
+                    btn_final_b = gr.Button("Execute Vault (Local) ➔", variant="secondary")
+                
+                with gr.Column(scale=1):
+                    out_final_b = gr.Textbox(label="Output: Reconstructed Document", lines=25)
+
+            btn_b.click(kernel_cloud_expert, inputs=inp_b, outputs=out_proto_b)
+            btn_final_b.click(kernel_vault, inputs=[out_proto_b, inp_db_b], outputs=out_final_b)
 
-        # --- TAB 2: REPORT ---
-        with gr.TabItem("📄 Technical Analysis"):
+        # --- TAB 3 ---
+        with gr.TabItem("📄 Technical Report"):
             gr.Markdown(report_md)
 
 demo.launch()