Create Medical_Mamba.py

Medical_Mamba.py

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+"""
+MEDICAL MAMBA v9.0 - PROPRIOCEPTIVE MEDICAL AI
+Offline/Online • Patient/Clinician • Probe-Steered
+
+Author: Logan Matthew Napolitano
+Proprioceptive AI, Inc. - February 2026
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import os
+import re
+import requests
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from ddgs import DDGS
+
+class C:
+    RESET = '\033[0m'
+    BOLD = '\033[1m'
+    DIM = '\033[2m'
+    RED = '\033[91m'
+    GREEN = '\033[92m'
+    YELLOW = '\033[93m'
+    CYAN = '\033[96m'
+    MAGENTA = '\033[95m'
+    BLUE = '\033[94m'
+    WHITE = '\033[97m'
+
+ESI_LEVELS = {
+    1: "Immediate - Life threatening",
+    2: "Emergent - High risk, severe pain",  
+    3: "Urgent - Stable but needs resources",
+    4: "Less Urgent - Stable, one resource",
+    5: "Non-Urgent - Stable"
+}
+
+RED_FLAGS = [
+    "chest pain", "difficulty breathing", "shortness of breath",
+    "severe bleeding", "coughing blood", "vomiting blood",
+    "sudden severe headache", "worst headache of life",
+    "facial droop", "arm weakness", "speech difficulty",
+    "loss of consciousness", "fainting", "unresponsive",
+    "severe abdominal pain", "rigid abdomen",
+    "suicidal", "homicidal", "self harm",
+    "allergic reaction", "anaphylaxis", "throat swelling",
+    "high fever", "fever over 103",
+    "trauma", "car accident", "fall from height", "assault",
+    "seizure", "convulsion",
+]
+
+class SymptomTracker:
+    def __init__(self, mode="patient"):
+        self.mode = mode
+        self.chief_complaint = None
+        self.symptoms = set()
+        self.red_flags_identified = []
+        self.conversation = []
+    
+    def check_red_flags(self, text):
+        text_lower = text.lower()
+        flags = [flag for flag in RED_FLAGS if flag in text_lower]
+        self.red_flags_identified.extend(flags)
+        return flags
+    
+    def extract_symptoms(self, text):
+        text_lower = text.lower()
+        if any(w in text_lower for w in ["leg", "thigh", "calf"]): self.symptoms.add("leg")
+        if any(w in text_lower for w in ["crush", "crushed", "fell on", "boulder", "trauma"]): 
+            self.symptoms.add("trauma")
+        if "severe" in text_lower or "10/10" in text_lower or "worst" in text_lower:
+            self.symptoms.add("severe_pain")
+        if "numb" in text_lower: self.symptoms.add("numbness")
+        if "weak" in text_lower: self.symptoms.add("weakness")
+    
+    def add_message(self, role, content):
+        self.conversation.append({"role": role, "content": content})
+        if role == "user":
+            self.check_red_flags(content)
+            self.extract_symptoms(content)
+            if self.chief_complaint is None:
+                self.chief_complaint = content
+    
+    def calculate_esi(self):
+        if self.red_flags_identified:
+            return 1 if any(f in ["chest pain", "difficulty breathing", "trauma"] for f in self.red_flags_identified) else 2
+        if "trauma" in self.symptoms: return 2
+        if "severe_pain" in self.symptoms: return 2
+        if len(self.symptoms) >= 3: return 3
+        return 4 if self.symptoms else 5
+
+# Probe architecture
+class FiberProjection(nn.Module):
+    def __init__(self, hidden_dim=4096, fiber_dim=16, n_layers=3):
+        super().__init__()
+        self.projections = nn.ModuleList([nn.Linear(hidden_dim, fiber_dim, bias=False) for _ in range(n_layers)])
+        self.layer_weights = nn.Parameter(torch.ones(n_layers) / n_layers)
+    
+    def forward(self, hidden_states, layer_indices):
+        projs = [self.projections[i](hidden_states[idx]) for i, idx in enumerate(layer_indices)]
+        stacked = torch.stack(projs, dim=0)
+        weights = F.softmax(self.layer_weights, dim=0).view(-1, 1, 1, 1)
+        return (weights * stacked).sum(dim=0)
+
+class ProbeHead(nn.Module):
+    def __init__(self, fiber_dim=16, hidden_dim=64):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(fiber_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1))
+    def forward(self, x):
+        return torch.sigmoid(self.net(x))
+
+class CognitiveProbe(nn.Module):
+    def __init__(self, hidden_dim=4096, fiber_dim=16, n_layers=3, head_hidden=64):
+        super().__init__()
+        self.fiber = FiberProjection(hidden_dim, fiber_dim, n_layers)
+        self.head = ProbeHead(fiber_dim, head_hidden)
+        self.layer_indices = [16, 32, 48]
+    
+    def forward(self, hidden_states):
+        return self.head(self.fiber(hidden_states, self.layer_indices))
+
+def load_probe(path, device):
+    if os.path.isdir(path):
+        for f in os.listdir(path):
+            if f.endswith('.pt'):
+                path = os.path.join(path, f)
+                break
+    ckpt = torch.load(path, map_location=device, weights_only=False)
+    probe = CognitiveProbe(hidden_dim=ckpt['hidden_dim'])
+    probe.layer_indices = ckpt['probe_layers']
+    probe.fiber.load_state_dict(ckpt['fiber_projection'])
+    probe.head.net.load_state_dict({k.replace('net.', ''): v for k, v in ckpt['head_state'].items()})
+    return probe.to(device).eval()
+
+SYSTEM_CLINICIAN = """You are a clinical decision support system with PROPRIOCEPTIVE SELF-AWARENESS.
+
+YOU ARE A CLINICIAN. ACT LIKE ONE:
+- Form clinical impressions based on the history
+- Identify concerning patterns and red flags
+- Give your differential thinking
+- Recommend specific actions
+
+When symptoms are concerning, state what you're worried about.
+Ask targeted follow-up questions to refine your assessment.
+
+WHEN YOU SEE [SELF-STATE calibration=0.6+]:
+- You're uncertain - acknowledge it
+- Say "I'm not certain but this could be..."
+
+DO NOT just tell them to "see a doctor" without giving your impression."""
+
+def main():
+    print(f"\n{C.CYAN}{'═'*60}{C.RESET}")
+    print(f"{C.CYAN}  MEDICAL MAMBA - PROPRIOCEPTIVE MEDICAL AI{C.RESET}")
+    print(f"{C.CYAN}{'═'*60}{C.RESET}\n")
+    
+    # Force CUDA, no auto device map
+    device = torch.device("cuda")
+    ROOT = "/home/programmer/Desktop/Claude_and_me/mamba7b_cognitive_output"
+    THRESHOLDS = {'depth': 0.65, 'specificity': 0.65, 'calibration': 0.55, 'coherence': 0.65, 'focus': 0.65}
+    
+    print(f"{C.WHITE}Loading Falcon-Mamba-7B...{C.RESET}")
+    tokenizer = AutoTokenizer.from_pretrained('tiiuae/falcon-mamba-7b-instruct', trust_remote_code=True)
+    
+    # Load directly to CUDA - no device_map='auto'
+    model = AutoModelForCausalLM.from_pretrained(
+        'tiiuae/falcon-mamba-7b-instruct',
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True
+    ).to(device)
+    model.eval()
+    print(f"{C.GREEN}✓ Model loaded on {device}{C.RESET}")
+    
+    print(f"{C.WHITE}Loading probes...{C.RESET}")
+    probes = {}
+    for name, ckpt in [('depth', 'ckpt_1000'), ('specificity', 'ckpt_1000'), ('calibration', 'ckpt_1500'), ('coherence', 'ckpt_1500'), ('focus', 'ckpt_1500')]:
+        path = os.path.join(ROOT, name, ckpt)
+        if os.path.exists(path):
+            probes[name] = load_probe(path, device)
+            print(f"  {C.GREEN}✓{C.RESET} {name}")
+    
+    tracker = SymptomTracker(mode="clinician")
+    print(f"\n{C.GREEN}Ready! Type 'quit' to exit.{C.RESET}\n")
+    
+    while True:
+        try:
+            user_input = input(f"{C.CYAN}HCP:{C.RESET} ").strip()
+            if not user_input:
+                continue
+            if user_input.lower() in ['quit', 'exit', 'q']:
+                break
+            
+            # Check red flags
+            flags = tracker.check_red_flags(user_input)
+            if flags:
+                print(f"\n{C.RED}⚠️  RED FLAG: {', '.join(flags)}{C.RESET}\n")
+            
+            tracker.add_message("user", user_input)
+            
+            # Build messages
+            messages = [
+                {"role": "system", "content": SYSTEM_CLINICIAN},
+                {"role": "user", "content": user_input}
+            ]
+            
+            # Add conversation history
+            if len(tracker.conversation) > 2:
+                messages = [{"role": "system", "content": SYSTEM_CLINICIAN}]
+                for msg in tracker.conversation[-6:]:
+                    messages.append(msg)
+            
+            prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            inputs = tokenizer(prompt, return_tensors='pt').to(device)
+            generated = inputs.input_ids.clone()
+            
+            scores_history = {n: [] for n in probes}
+            state_injections = 0
+            generated_text = ""
+            
+            print(f"\n{C.GREEN}Mamba:{C.RESET} ", end="", flush=True)
+            
+            with torch.no_grad():
+                for step in range(500):
+                    outputs = model(generated, output_hidden_states=True, return_dict=True)
+                    hidden_states = list(outputs.hidden_states)
+                    
+                    # Read probe scores
+                    current = {}
+                    for name, probe in probes.items():
+                        score = probe(hidden_states)[0, -1].item()
+                        current[name] = score
+                        scores_history[name].append(score)
+                    
+                    problems = [n for n, s in current.items() if s > THRESHOLDS[n]]
+                    
+                    # Self-state injection every 20 tokens if problems
+                    if problems and step > 0 and step % 20 == 0:
+                        state_parts = [f"{n}={current[n]:.2f}" for n in problems]
+                        state_msg = f" [SELF-STATE: {' '.join(state_parts)}] "
+                        state_tokens = tokenizer.encode(state_msg, add_special_tokens=False)
+                        for st in state_tokens:
+                            generated = torch.cat([generated, torch.tensor([[st]], device=device)], dim=1)
+                        state_injections += 1
+                        print(f"{C.MAGENTA}{state_msg}{C.RESET}", end="", flush=True)
+                    
+                    # Temperature based on probe state
+                    temp = 0.35 if problems else 0.6
+                    
+                    logits = outputs.logits[:, -1, :] / temp
+                    next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+                    
+                    token_str = tokenizer.decode(next_token[0])
+                    generated_text += token_str
+                    
+                    # Color code output
+                    if 'calibration' in problems:
+                        print(f"{C.RED}{token_str}{C.RESET}", end="", flush=True)
+                    elif problems:
+                        print(f"{C.YELLOW}{token_str}{C.RESET}", end="", flush=True)
+                    else:
+                        print(token_str, end="", flush=True)
+                    
+                    generated = torch.cat([generated, next_token], dim=1)
+                    if next_token.item() == tokenizer.eos_token_id:
+                        break
+            
+            tracker.add_message("assistant", generated_text.strip())
+            
+            # Summary
+            print(f"\n\n{C.DIM}{'─'*60}{C.RESET}")
+            esi = tracker.calculate_esi()
+            esi_color = C.RED if esi <= 2 else C.YELLOW if esi == 3 else C.GREEN
+            print(f"{C.WHITE}ESI:{C.RESET} {esi_color}Level {esi}{C.RESET} | {C.MAGENTA}SELF-STATES:{C.RESET} {state_injections} | ", end="")
+            
+            for n in ['calibration', 'coherence']:
+                if scores_history.get(n):
+                    avg = sum(scores_history[n]) / len(scores_history[n])
+                    c = C.GREEN if avg < 0.4 else C.YELLOW if avg < 0.55 else C.RED
+                    print(f"{n}:{c}{avg:.2f}{C.RESET} ", end="")
+            print(f"\n{C.DIM}{'─'*60}{C.RESET}\n")
+            
+        except KeyboardInterrupt:
+            print(f"\n{C.YELLOW}(Use 'quit' to exit){C.RESET}")
+            continue
+        except Exception as e:
+            print(f"\n{C.RED}Error: {e}{C.RESET}")
+            continue
+    
+    print(f"\n{C.CYAN}Session ended.{C.RESET}\n")
+
+if __name__ == "__main__":
+    main()