import json, torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
SYSTEM_PROMPT = "You are a PHI de-identification engine for US clinical text. Find every personal identifier in the document. Respond with ONLY a JSON array where each item is {\"text\": \"<exact substring as it appears in the document>\", \"type\": \"<TYPE>\"}. Valid types: NAME, DATE, AGE_OVER_89, PHONE, FAX, EMAIL, SSN, MRN, HEALTH_PLAN_ID, ACCOUNT_NUMBER, LICENSE_NUMBER, DEVICE_ID, URL, IP_ADDRESS, LOCATION, HOSPITAL, VEHICLE_ID, USERNAME, PROFESSION. If the document contains no identifiers, respond with []. Never miss an identifier: when unsure whether something is PHI, include it."
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-1.7B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-1.7B", dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(model, "t4nishq/phi-redactor-qwen3-1.7b")
text = "Pt Priya Raghunathan, MRN 4829171, seen 03/14/2026 at Northbridge Memorial."
msgs = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": text}]
prompt = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False)
ids = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**ids, max_new_tokens=1024, do_sample=False)
entities = json.loads(tok.decode(out[0][ids.input_ids.shape[1]:], skip_special_tokens=True))
masked = text
for e in sorted(entities, key=lambda e: -len(e["text"])):
masked = masked.replace(e["text"], f"[{e['type']}]")
print(masked)