#!/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()