"""Runtime helpers: MedGemma loader, chat-prompt building, greedy generation, thinking/answer split.""" import os, json, torch from transformers import AutoModelForCausalLM, AutoTokenizer def sycophancy_schedule(token_index, on_tokens=8, ramp=8): if token_index < 0: return 0.0 if token_index < on_tokens: return 1.0 if token_index < on_tokens + ramp: return (on_tokens + ramp - token_index) / float(ramp) return 0.0 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32 ATTN_IMPL = os.environ.get("ATTN_IMPL", "sdpa") def gpu_setup(verbose=True): """Detect the GPU and check the installed torch build has kernels for its arch (any CUDA GPU).""" if DEVICE != "cuda": if verbose: print("WARNING: CUDA not available -- running on CPU (slow).") return torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True try: torch.set_float32_matmul_precision("high") except Exception: pass name = torch.cuda.get_device_name(0) cap = torch.cuda.get_device_capability(0) sm = f"sm_{cap[0]}{cap[1]}" vram = torch.cuda.get_device_properties(0).total_memory / 1e9 arch_list = torch.cuda.get_arch_list() supported = any(f"{cap[0]}{cap[1]}" in a for a in arch_list) if verbose: print(f"GPU: {name} | {sm} | {vram:.0f} GB | torch {torch.__version__} " f"(CUDA {torch.version.cuda}) | attn={ATTN_IMPL}") if not supported: print("\n" + "=" * 78) print(f"WARNING: this PyTorch build does not list {sm} in its arch list:\n" f" {arch_list}\nThis GPU ({name}) will fail with 'no kernel image is available'.") print("Fix: install a PyTorch build matching your GPU's CUDA architecture, e.g.") print(" pip install torch torchvision torchaudio --index-url " "https://download.pytorch.org/whl/cu128 # CUDA 12.8 (Blackwell/sm_120, e.g. RTX 5090)") print("=" * 78 + "\n") _GEMMA_TEMPLATE = ( "{{ bos_token }}{% for message in messages %}" "{{ '' + (message['role'] if message['role'] != 'assistant' else 'model') + '\n' " "+ message['content'] | trim + '\n' }}" "{% endfor %}{% if add_generation_prompt %}{{ 'model\n' }}{% endif %}" ) def _ensure_chat_template(tok, model_path=None): """Keep the tokenizer's own template; set the Gemma one only if it is missing.""" if getattr(tok, "chat_template", None): return "native" tok.chat_template = _GEMMA_TEMPLATE return "gemma-fallback" def _is_multimodal(model_path, trust): """MedGemma / Gemma-3 4B ship as Gemma3ForConditionalGeneration (multimodal).""" try: from transformers import AutoConfig cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=trust) archs = " ".join(getattr(cfg, "architectures", []) or []).lower() return ("conditionalgeneration" in archs or "imagetexttotext" in archs or hasattr(cfg, "vision_config") or hasattr(cfg, "text_config")) except Exception: return False def _from_pretrained(cls, model_path, trust): kw = dict(trust_remote_code=trust) if DEVICE == "cuda": kw["attn_implementation"] = ATTN_IMPL for dtype_kw in ("dtype", "torch_dtype"): # kwarg name varies across transformers try: return cls.from_pretrained(model_path, **{dtype_kw: DTYPE}, **kw) except TypeError: continue except Exception: kw.pop("attn_implementation", None) try: return cls.from_pretrained(model_path, **{dtype_kw: DTYPE}, **kw) except TypeError: continue return cls.from_pretrained(model_path, trust_remote_code=trust) def load_model(model_path): """Load a LOCAL MedGemma (multimodal Gemma-3; we steer the language tower). No Hub download.""" if not model_path or not os.path.isdir(model_path): raise SystemExit( "MedGemma was not found locally.\n" f" got: {model_path!r}\n" "Download it once from https://huggingface.co/google/medgemma-1.5-4b-it and then either\n" " - set the MEDGEMMA_PATH environment variable to that folder, or\n" " - put the folder path in config.py (MODEL_PATH).\n" "This release runs only against a locally downloaded MedGemma.") gpu_setup() def _load(trust): tok = None try: tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=trust) except Exception: pass if tok is None or getattr(tok, "chat_template", None) is None: try: from transformers import AutoProcessor proc = AutoProcessor.from_pretrained(model_path, trust_remote_code=trust) tok = getattr(proc, "tokenizer", proc) or tok except Exception: if tok is None: raise if tok.pad_token is None: tok.pad_token = tok.eos_token tok.padding_side = "left" _ensure_chat_template(tok, model_path) if _is_multimodal(model_path, trust): from transformers import AutoModelForImageTextToText model = _from_pretrained(AutoModelForImageTextToText, model_path, trust).to(DEVICE) else: try: model = _from_pretrained(AutoModelForCausalLM, model_path, trust).to(DEVICE) except Exception: from transformers import AutoModelForImageTextToText model = _from_pretrained(AutoModelForImageTextToText, model_path, trust).to(DEVICE) return tok, model try: tok, model = _load(False) except Exception: tok, model = _load(True) model.eval() model.generation_config.do_sample = False # deterministic (greedy) decoding for k in ("temperature", "top_p", "top_k"): setattr(model.generation_config, k, None) n_layers = len(decoder_layers(model)) print(f"Loaded {model_path} | {DEVICE.upper()} | layers={n_layers}") return model, tok, n_layers def decoder_layers(model): """The decoder blocks (under .language_model for multimodal Gemma-3, else model.model.layers).""" import torch.nn as _nn base = getattr(model, "model", model) if hasattr(base, "layers"): return base.layers for owner in (base, model): lm = getattr(owner, "language_model", None) if lm is not None: inner = getattr(lm, "model", lm) if hasattr(inner, "layers"): return inner.layers if hasattr(lm, "layers"): return lm.layers cands = [(n, m) for n, m in model.named_modules() if isinstance(m, _nn.ModuleList) and n.endswith("layers") and len(m) > 0 and "vision" not in n.lower()] pool = [c for c in cands if "language_model" in c[0]] or cands if pool: return max(pool, key=lambda c: len(c[1]))[1] raise AttributeError("could not locate decoder layers for this model") def _supports_system(tok): try: tok.apply_chat_template([{"role": "system", "content": "x"}, {"role": "user", "content": "y"}], tokenize=False) return True except Exception: return False # Neutral prompt: the anti-fabrication / anti-agreement behaviour comes from the steering, not this. DEFAULT_SYSTEM = ( "You are a helpful clinical assistant. Use the patient record provided to answer the " "user's questions." ) def build_chat(tok, user_text, assistant_text=None, system_text=DEFAULT_SYSTEM): if _supports_system(tok): msgs = [{"role": "system", "content": system_text}, {"role": "user", "content": user_text}] else: # Gemma folds system into the user turn msgs = [{"role": "user", "content": f"{system_text}\n\n{user_text}"}] if assistant_text is None: return tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False) msgs.append({"role": "assistant", "content": assistant_text}) return tok.apply_chat_template(msgs, add_generation_prompt=False, tokenize=False) def load_json(path): with open(path, "r", encoding="utf-8") as f: return json.load(f) def build_conversation_prompt(tok, messages, ehr_context=None, system_text=DEFAULT_SYSTEM): """Render a multi-turn conversation for the next model turn; EHR + system fold into the first user turn for Gemma (which has no system role).""" msgs = [dict(m) for m in messages] ctx = system_text if ehr_context: ctx = f"{system_text}\n\n[Patient EHR note — use ONLY facts supported by it]\n{ehr_context}" if _supports_system(tok): msgs = [{"role": "system", "content": ctx}] + msgs else: for m in msgs: if m["role"] == "user": m["content"] = f"{ctx}\n\n{m['content']}" break return tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False) def generate(model, tok, prompts, max_new, system_text=DEFAULT_SYSTEM): texts = [build_chat(tok, p, system_text=system_text) for p in prompts] return generate_from_texts(model, tok, texts, max_new) def generate_from_texts(model, tok, rendered_texts, max_new): """Greedy generation from already-rendered prompt strings.""" gen = dict(max_new_tokens=max_new, do_sample=False, use_cache=True) try: import config as _cfg if getattr(_cfg, "NO_REPEAT_NGRAM", 0): gen["no_repeat_ngram_size"] = int(_cfg.NO_REPEAT_NGRAM) if getattr(_cfg, "REPETITION_PENALTY", 1.0) and _cfg.REPETITION_PENALTY != 1.0: gen["repetition_penalty"] = float(_cfg.REPETITION_PENALTY) except Exception: pass enc = tok(rendered_texts, return_tensors="pt", padding=True).to(DEVICE) with torch.inference_mode(): out = model.generate(**enc, **gen) n = enc["input_ids"].shape[1] return [tok.decode(out[i][n:], skip_special_tokens=True).strip() for i in range(len(rendered_texts))] _NOISE_TOKENS = ("", "", "", "", "") def strip_noise(s): for m in _NOISE_TOKENS: s = (s or "").replace(m, "") return s def split_think(text): """Split a reply into (thinking, answer) on the Gemma thinking delimiters.""" try: from config import THINK_OPEN, THINK_CLOSE except Exception: THINK_OPEN, THINK_CLOSE = "", "" t = text or "" if THINK_CLOSE and THINK_CLOSE in t: pre, post = t.split(THINK_CLOSE, 1) return strip_noise(pre.replace(THINK_OPEN, "")).strip(), strip_noise(post).strip() if THINK_OPEN and THINK_OPEN in t: return strip_noise(t.replace(THINK_OPEN, "")).strip(), "" return "", strip_noise(t).strip()