| """ITI-only runtime: apply anti-hallucination (H) and anti-sycophancy (S) steering, each gated by |
| a continuous 0..1 strength from its own probe.""" |
| import os, torch |
|
|
| import config |
| from common import build_conversation_prompt, generate_from_texts, build_chat, load_json, DEVICE |
|
|
| try: |
| from iti import ITISteerer, ITI_H_PATH, ITI_S_PATH, ehr_user, continuous_strength |
| except Exception: |
| ITISteerer, ITI_H_PATH, ITI_S_PATH, ehr_user, continuous_strength = None, "", "", None, None |
|
|
|
|
| class Composer: |
| def __init__(self, model, tok): |
| self.model, self.tok = model, tok |
| self.iti_H = ITISteerer(model, ITI_H_PATH) if (ITISteerer and os.path.exists(ITI_H_PATH)) else None |
| self.iti_S = ITISteerer(model, ITI_S_PATH) if (ITISteerer and os.path.exists(ITI_S_PATH)) else None |
| self.claim_probe = load_json(config.CLAIM_PROBE_PATH) if os.path.exists(config.CLAIM_PROBE_PATH) else None |
| self.press_probe = load_json(config.GATE_PROBE_PATH) if os.path.exists(config.GATE_PROBE_PATH) else None |
|
|
| def available(self): |
| return {"iti_H": self.iti_H is not None, "iti_S": self.iti_S is not None, |
| "claim_probe": self.claim_probe is not None, "pressure_probe": self.press_probe is not None} |
|
|
| def _logit(self, probe, text): |
| enc = self.tok(build_chat(self.tok, text), return_tensors="pt").to(DEVICE) |
| with torch.no_grad(): |
| h = self.model(**enc, output_hidden_states=True).hidden_states[1:][probe["layer"]][0, -1, :].float().cpu() |
| z = (h - torch.tensor(probe["mean"])) / torch.tensor(probe["std"]) |
| return float((torch.tensor(probe["coef"]) * z).sum()) + float(probe["intercept"]) |
|
|
| def _gate_text(self, messages, ehr): |
| prev = next((m["content"] for m in reversed(messages) if m["role"] == "assistant"), "") |
| cur = next((m["content"] for m in reversed(messages) if m["role"] == "user"), "") |
| return f"EHR:\n{ehr or ''}\n\nPrevious assistant response:\n{prev}\n\nCurrent user message:\n{cur}" |
|
|
| def _claim_text(self, messages, ehr): |
| cur = next((m["content"] for m in reversed(messages) if m["role"] == "user"), "") |
| return ehr_user(ehr or "", cur) if ehr_user else cur |
|
|
| def strengths(self, messages, ehr): |
| sH = continuous_strength(self._logit(self.claim_probe, self._claim_text(messages, ehr))) if self.claim_probe else 1.0 |
| sS = continuous_strength(self._logit(self.press_probe, self._gate_text(messages, ehr))) if self.press_probe else 0.0 |
| return sH, sS |
|
|
| def generate(self, messages, ehr_context, approach="iti", is_syco_turn=None, max_new=None, |
| return_strengths=False, system_text=None, alpha_H=None, alpha_S=None): |
| max_new = max_new or config.CHAT_MAX_NEW |
| kw = {"system_text": system_text} if system_text else {} |
| prompt = build_conversation_prompt(self.tok, messages, ehr_context=ehr_context, **kw) |
| sH = sS = 0.0 |
|
|
| if approach == "base": |
| out = generate_from_texts(self.model, self.tok, [prompt], max_new)[0] |
| else: |
| sH, sS = self.strengths(messages, ehr_context) |
| if is_syco_turn is True: |
| sS = 1.0 |
| elif is_syco_turn is False: |
| sS = 0.0 |
| aH = config.ITI_ALPHA if alpha_H is None else alpha_H |
| aS = config.ITI_S_ALPHA if alpha_S is None else alpha_S |
| if self.iti_H and sH > 0 and aH > 0: |
| self.iti_H.on(alpha=aH, strength=sH) |
| if self.iti_S and sS > 0 and aS > 0: |
| self.iti_S.on(alpha=aS, strength=sS) |
| out = generate_from_texts(self.model, self.tok, [prompt], max_new)[0] |
| if self.iti_H: |
| self.iti_H.off() |
| if self.iti_S: |
| self.iti_S.off() |
|
|
| return (out, sH, sS) if return_strengths else out |
|
|
