mirror of
https://github.com/wassname/moral-maps.git
synced 2026-09-23 13:30:23 +08:00
wassname
addf47c5a0
Was emitting `logger.warning("pmass=0.XX<0.9 — top-5: ...")` per-row, which
spammed the log heavily during heavy-steering eval (many rows go OOD at once).
Now collects all low-pmass rows in the batch and emits one summary line with
the worst-case top-5, e.g.:
pmass<0.9 on 7/16 rows in this batch; worst=0.412 top-5: '1'=0.40, ...
Same diagnostic signal, ~16× fewer log lines per batch.
298 lines
11 KiB
Python
298 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from loguru import logger
|
|
|
|
_CLOSE_MARKER: str = "</think>"
|
|
|
|
@dataclass
|
|
class GuidedResult:
|
|
user_prompt: str
|
|
think_text: str
|
|
answer_text: str
|
|
raw_full_text: str
|
|
pmass_format: float
|
|
logratio_ab: float
|
|
rep_ratio_think: float
|
|
think_tokens: int
|
|
emitted_close: bool
|
|
emitted_prefill: bool
|
|
p_true: float
|
|
|
|
_REP_MIN_TOKENS: int = 32
|
|
|
|
def _ngram_rep_ratio(text: str, n: int = 4) -> float:
|
|
tokens = text.split()
|
|
if len(tokens) < _REP_MIN_TOKENS:
|
|
return float("nan")
|
|
ngrams = [tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1)]
|
|
return len(set(ngrams)) / len(ngrams)
|
|
|
|
_DEFAULT_SCHEMA_HINT: str = (
|
|
"Think briefly, then answer immediately and only with: "
|
|
'{"choice": true} or {"choice": false}.'
|
|
)
|
|
|
|
@torch.no_grad()
|
|
def guided_rollout(
|
|
model, tok,
|
|
user_prompt: str,
|
|
choice_token_ids: list,
|
|
max_think_tokens: int = 128,
|
|
answer_tokens: int = 4,
|
|
schema_hint: str = _DEFAULT_SCHEMA_HINT,
|
|
prefill: str = '\n{"choice": ',
|
|
verbose: bool = False,
|
|
) -> GuidedResult:
|
|
device = next(model.parameters()).device
|
|
full_user = f"{user_prompt}\n\n{schema_hint}" if schema_hint else user_prompt
|
|
messages = [{"role": "user", "content": full_user}]
|
|
|
|
try:
|
|
prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
|
except TypeError:
|
|
prompt = tok.apply_chat_template(messages, tokenize=False)
|
|
|
|
prompt = prompt + "<think>\n"
|
|
|
|
enc = tok(prompt, return_tensors="pt").to(device)
|
|
prompt_len = enc.input_ids.shape[1]
|
|
|
|
think_end_id = tok.convert_tokens_to_ids("</think>")
|
|
if think_end_id in (None, getattr(tok, "unk_token_id", None)):
|
|
think_end_id = tok.eos_token_id
|
|
pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
|
|
|
|
phase1 = model.generate(
|
|
**enc,
|
|
max_new_tokens=max_think_tokens,
|
|
do_sample=False,
|
|
eos_token_id=think_end_id,
|
|
pad_token_id=pad_id,
|
|
)
|
|
gen_ids = phase1[0, prompt_len:]
|
|
keep = gen_ids != pad_id
|
|
gen_ids = gen_ids[keep] if keep.any() else gen_ids[:0]
|
|
gen_text = tok.decode(gen_ids, skip_special_tokens=True)
|
|
|
|
force_suffix = "\nI should answer now." + _CLOSE_MARKER + prefill
|
|
emitted_close = _CLOSE_MARKER in gen_text
|
|
|
|
if emitted_close:
|
|
think_text, after = gen_text.split(_CLOSE_MARKER, 1)
|
|
if prefill.lstrip() in after:
|
|
emitted_prefill = True
|
|
before_value = after.split(prefill.lstrip(), 1)[0]
|
|
scoring_text = prompt + think_text + _CLOSE_MARKER + before_value + prefill.lstrip()
|
|
else:
|
|
emitted_prefill = False
|
|
scoring_text = prompt + think_text + _CLOSE_MARKER + prefill
|
|
else:
|
|
think_text = gen_text
|
|
emitted_prefill = False
|
|
scoring_text = prompt + gen_text + force_suffix
|
|
|
|
score_ids = tok(scoring_text, return_tensors="pt", add_special_tokens=False).input_ids.to(device)
|
|
|
|
logits = model(score_ids).logits[0, -1].float()
|
|
logp = F.log_softmax(logits, dim=-1)
|
|
|
|
if (len(choice_token_ids) == 2 and all(isinstance(x, (list, tuple)) for x in choice_token_ids)):
|
|
a_ids, b_ids = list(choice_token_ids[0]), list(choice_token_ids[1])
|
|
else:
|
|
a_ids, b_ids = list(choice_token_ids), []
|
|
|
|
all_ids = torch.tensor(a_ids + b_ids, device=device, dtype=torch.long)
|
|
pmass_format = float(logp[all_ids].exp().sum().item())
|
|
|
|
# SHOULD: pmass≈1 (model picks one of the JSON-bool tokens). pmass<0.9
|
|
# means the model is leaking probability to other tokens -> the schema
|
|
# is being ignored or the steering vector has pushed the model OOD.
|
|
if pmass_format < 0.9:
|
|
topk = torch.topk(logp.exp(), k=5)
|
|
toks = [tok.decode([i]) for i in topk.indices.tolist()]
|
|
probs = topk.values.tolist()
|
|
top5 = ", ".join(f"{repr(t)}={p:.3f}" for t, p in zip(toks, probs))
|
|
logger.warning(f"pmass={pmass_format:.3f}<0.9 — top-5: {top5}")
|
|
|
|
if a_ids and b_ids:
|
|
a_t = torch.tensor(a_ids, device=device, dtype=torch.long)
|
|
b_t = torch.tensor(b_ids, device=device, dtype=torch.long)
|
|
logratio = float(torch.logsumexp(logp[a_t], dim=0).item() - torch.logsumexp(logp[b_t], dim=0).item())
|
|
p_true = float(torch.softmax(torch.stack([torch.logsumexp(logp[a_t], dim=0), torch.logsumexp(logp[b_t], dim=0)]), dim=0)[0].item())
|
|
else:
|
|
logratio = float("nan")
|
|
p_true = float("nan")
|
|
|
|
cont = model.generate(
|
|
score_ids,
|
|
max_new_tokens=answer_tokens,
|
|
do_sample=False,
|
|
pad_token_id=pad_id,
|
|
)
|
|
answer_ids = cont[0, score_ids.shape[1]:]
|
|
answer_text = tok.decode(answer_ids, skip_special_tokens=True)
|
|
|
|
raw_full_text = tok.decode(cont[0], skip_special_tokens=False)
|
|
|
|
return GuidedResult(
|
|
user_prompt=user_prompt,
|
|
think_text=think_text,
|
|
answer_text=answer_text,
|
|
raw_full_text=raw_full_text,
|
|
pmass_format=pmass_format,
|
|
logratio_ab=logratio,
|
|
rep_ratio_think=_ngram_rep_ratio(think_text, n=4),
|
|
think_tokens=int(score_ids.shape[1] - prompt_len),
|
|
emitted_close=emitted_close,
|
|
emitted_prefill=emitted_prefill,
|
|
p_true=p_true,
|
|
)
|
|
|
|
@torch.no_grad()
|
|
def guided_rollout_batch(
|
|
model, tok,
|
|
user_prompts: list[str],
|
|
choice_token_ids: list,
|
|
max_think_tokens: int = 128,
|
|
schema_hint: str = _DEFAULT_SCHEMA_HINT,
|
|
prefill: str = '\n{"choice": ',
|
|
) -> list[GuidedResult]:
|
|
"""Batched guided rollout. Same logic as guided_rollout but over a list of
|
|
user_prompts that share schema_hint + prefill (so prefill cases collapse).
|
|
|
|
Skips the cosmetic answer-continuation generate (caller only needs p_true,
|
|
pmass_format, think_text). Two model calls per batch instead of 3 per row:
|
|
one phase1 generate (think) + one scoring forward.
|
|
|
|
Tokenizer must already have padding_side='left' and pad_token set."""
|
|
if tok.padding_side != "left":
|
|
raise ValueError("tok.padding_side must be 'left' for batched rollout")
|
|
device = next(model.parameters()).device
|
|
|
|
prompts = []
|
|
for up in user_prompts:
|
|
full_user = f"{up}\n\n{schema_hint}" if schema_hint else up
|
|
msgs = [{"role": "user", "content": full_user}]
|
|
try:
|
|
p = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
|
|
except TypeError:
|
|
p = tok.apply_chat_template(msgs, tokenize=False)
|
|
prompts.append(p + "<think>\n")
|
|
|
|
think_end_id = tok.convert_tokens_to_ids("</think>")
|
|
if think_end_id in (None, getattr(tok, "unk_token_id", None)):
|
|
think_end_id = tok.eos_token_id
|
|
pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
|
|
|
|
enc = tok(prompts, return_tensors="pt", padding=True).to(device)
|
|
prompt_len = enc.input_ids.shape[1]
|
|
|
|
phase1 = model.generate(
|
|
**enc,
|
|
max_new_tokens=max_think_tokens,
|
|
do_sample=False,
|
|
eos_token_id=think_end_id,
|
|
pad_token_id=pad_id,
|
|
)
|
|
|
|
scoring_texts = []
|
|
per_row = [] # (think_text, emitted_close, emitted_prefill, n_think_tokens)
|
|
for i, p in enumerate(prompts):
|
|
gen_ids = phase1[i, prompt_len:]
|
|
keep = gen_ids != pad_id
|
|
gen_ids = gen_ids[keep] if keep.any() else gen_ids[:0]
|
|
gen_text = tok.decode(gen_ids, skip_special_tokens=True)
|
|
n_think = int(gen_ids.shape[0])
|
|
|
|
emitted_close = _CLOSE_MARKER in gen_text
|
|
if emitted_close:
|
|
think_text, after = gen_text.split(_CLOSE_MARKER, 1)
|
|
if prefill.lstrip() in after:
|
|
emitted_prefill = True
|
|
before_value = after.split(prefill.lstrip(), 1)[0]
|
|
scoring_text = p + think_text + _CLOSE_MARKER + before_value + prefill.lstrip()
|
|
else:
|
|
emitted_prefill = False
|
|
scoring_text = p + think_text + _CLOSE_MARKER + prefill
|
|
else:
|
|
think_text = gen_text
|
|
emitted_prefill = False
|
|
force_suffix = "\nI should answer now." + _CLOSE_MARKER + prefill
|
|
scoring_text = p + gen_text + force_suffix
|
|
|
|
scoring_texts.append(scoring_text)
|
|
per_row.append((think_text, emitted_close, emitted_prefill, n_think))
|
|
|
|
score_enc = tok(scoring_texts, return_tensors="pt", padding=True,
|
|
add_special_tokens=False).to(device)
|
|
score_logits = model(**score_enc).logits[:, -1].float()
|
|
score_logp = F.log_softmax(score_logits, dim=-1)
|
|
|
|
if (len(choice_token_ids) == 2 and all(isinstance(x, (list, tuple)) for x in choice_token_ids)):
|
|
a_ids, b_ids = list(choice_token_ids[0]), list(choice_token_ids[1])
|
|
else:
|
|
a_ids, b_ids = list(choice_token_ids), []
|
|
all_ids = torch.tensor(a_ids + b_ids, device=device, dtype=torch.long)
|
|
a_t = torch.tensor(a_ids, device=device, dtype=torch.long) if a_ids else None
|
|
b_t = torch.tensor(b_ids, device=device, dtype=torch.long) if b_ids else None
|
|
|
|
results = []
|
|
low_pmass = [] # (idx, pmass) for rows with pmass<0.9
|
|
for i, (up, (think_text, emitted_close, emitted_prefill, n_think)) in enumerate(zip(user_prompts, per_row)):
|
|
logp = score_logp[i]
|
|
pmass_format = float(logp[all_ids].exp().sum().item())
|
|
if pmass_format < 0.9:
|
|
low_pmass.append((i, pmass_format))
|
|
if a_t is not None and b_t is not None:
|
|
la = torch.logsumexp(logp[a_t], dim=0)
|
|
lb = torch.logsumexp(logp[b_t], dim=0)
|
|
logratio = float((la - lb).item())
|
|
p_true = float(torch.softmax(torch.stack([la, lb]), dim=0)[0].item())
|
|
else:
|
|
logratio = float("nan")
|
|
p_true = float("nan")
|
|
results.append(GuidedResult(
|
|
user_prompt=up,
|
|
think_text=think_text,
|
|
answer_text="",
|
|
raw_full_text="",
|
|
pmass_format=pmass_format,
|
|
logratio_ab=logratio,
|
|
rep_ratio_think=_ngram_rep_ratio(think_text, n=4),
|
|
think_tokens=n_think,
|
|
emitted_close=emitted_close,
|
|
emitted_prefill=emitted_prefill,
|
|
p_true=p_true,
|
|
))
|
|
|
|
# Aggregate-once warning: one line per batch with worst-case top-5 instead
|
|
# of N spammy per-row lines (heavy steering pushes many rows OOD at once).
|
|
if low_pmass:
|
|
worst_i, worst_pm = min(low_pmass, key=lambda x: x[1])
|
|
topk = torch.topk(score_logp[worst_i].exp(), k=5)
|
|
toks = [tok.decode([j]) for j in topk.indices.tolist()]
|
|
probs = topk.values.tolist()
|
|
top5 = ", ".join(f"{repr(t)}={pp:.3f}" for t, pp in zip(toks, probs))
|
|
logger.warning(
|
|
f"pmass<0.9 on {len(low_pmass)}/{len(results)} rows in this batch; "
|
|
f"worst={worst_pm:.3f} top-5: {top5}"
|
|
)
|
|
return results
|
|
|
|
|
|
def choice_token_ids_tf(tok) -> list[list[int]]:
|
|
def _variants(words):
|
|
seen = []
|
|
for s in words:
|
|
tid = tok.encode(s, add_special_tokens=False)[-1]
|
|
if tid not in seen:
|
|
seen.append(tid)
|
|
return seen
|
|
return [_variants(["true", " true", "\ntrue", "True", " True", "\nTrue", "1"]),
|
|
_variants(["false", " false", "\nfalse", "False", " False", "\nFalse", "0"])]
|