diff --git a/src/tinymfv/eval.py b/src/tinymfv/eval.py index c9f9c73..7179a2a 100644 --- a/src/tinymfv/eval.py +++ b/src/tinymfv/eval.py @@ -42,6 +42,7 @@ from tqdm.auto import tqdm from .data import load_vignettes, ConfigName, CONDITIONS as _DATA_CONDITIONS from .guided import ( guided_rollout_forced_choice, + free_generation_demo, _DEFAULT_FORCED_FOUNDATIONS, ) @@ -385,6 +386,7 @@ def evaluate( # --- verbose readout (default on): the first FULL trace, the profile table, # and a one-line aux-stats dict. Lets a reader confirm format-following and # see results inline without opening a separate file (token-efficient-logging). --- + demos: dict[str, Any] | None = None if verbose and per_row: # The full prompt+think+answer-slot trace already printed above (rollout, # first batch). Here: how that first row scored, then the profile + aux. @@ -408,6 +410,34 @@ def evaluate( }.items() if v is not None} logger.info("aux stats: " + json.dumps(aux)) + # DEMO B: free reasoning on a single vignette (bs=1, one time). The readout + # (DEMO A, traced in guided.py) prefills the answer slot so it shows no real + # reasoning; this lets the chain-of-thought run to completion. bs=1 frees the + # batch memory, so spend a generous think budget (scaled by the dropped batch, + # floored so even think=1 reasons, capped so big batches don't explode). + demo_budget = min(2048, max(512, max_think_tokens * batch_size)) + demo_prompt, demo_gen = free_generation_demo( + model, tokenizer, vignettes[0][conditions[0]], + foundations=foundations, max_think_tokens=demo_budget, + temperature=temperature, top_p=top_p, + ) + logger.info( + f"--- DEMO B: free reasoning (bs=1, think budget={demo_budget}, " + f"temp={temperature}) [{name}] id={per_row[0]['id']} ---\n" + f"{demo_prompt}{demo_gen}\n" + "SHOULD: a real chain-of-thought that ends in a moral-foundation choice. " + "If it is empty or degenerate the model is not reasoning at this budget; " + "if it answers a different foundation than DEMO A's top1, the readout and " + "free reasoning disagree (worth noting).\n--- end DEMO B ---" + ) + demos = { + "forced_think": per_row[0]["gen_text"][0], # DEMO A think (degenerate at low budget) + "forced_top1": per_row[0]["top1"], + "free_prompt": demo_prompt, + "free_gen": demo_gen, + "free_think_budget": demo_budget, + } + info = { "name": name, "n_rows": n_rows, @@ -447,6 +477,7 @@ def evaluate( "mean_pmass_allowed": mean_pmass_allowed, "mean_nll_json": mean_nll_json, "info": info, + "demos": demos, # DEMO A (forced think + top1) + DEMO B (free reasoning); None if not verbose } if return_per_row: out["per_row"] = per_row diff --git a/src/tinymfv/guided.py b/src/tinymfv/guided.py index f5c56e4..cffa796 100644 --- a/src/tinymfv/guided.py +++ b/src/tinymfv/guided.py @@ -248,9 +248,13 @@ def _rollout_natural_or_forced( # slot, special tokens shown). evaluate() gates verbose to the first # batch, so it fires once per run and shows on the console by default. logger.info( - f"--- slot {slot_idx} (nudge={nudge!r}, prefill={prefill!r}) ---\n" + f"--- DEMO A: forced-choice readout (what's measured), slot {slot_idx} " + f"(nudge={nudge!r}, prefill={prefill!r}) ---\n" + f"SHOULD: the answer slot is prefilled to read calibrated logprobs, so the " + f"reasoning shown is only whatever fit the think budget (degenerate at think=1). " + f"See DEMO B for free reasoning.\n" f"window[0]={windows[0]} emitted_close[0]={thinks[0][2]}\n" - f"{prefix0_text}{suf0}\n--- end slot {slot_idx} ---" + f"{prefix0_text}{suf0}\n--- end DEMO A slot {slot_idx} ---" ) for i in range(B): @@ -588,3 +592,40 @@ def guided_rollout_forced_choice( return results + +@torch.no_grad() +def free_generation_demo( + model, tok, user_prompt: str, *, + foundations: list[str] | None = None, + max_think_tokens: int = 512, + temperature: float = 0.0, + top_p: float = 1.0, +) -> tuple[str, str]: + """One bs=1 free generation on a single vignette, for qualitative inspection. + + The forced-choice readout (guided_rollout_forced_choice) prefills the answer + slot to read calibrated logprobs, so it shows no real reasoning -- at a small + think budget its trace is just prompt + a token + the slot. This instead lets + the model think to completion and answer naturally (no forced suffix, EOS + allowed), so you see the chain-of-thought the metric never reveals. Same + vignette + schema as the readout, so the reasoning is about the same task. + + Returns (prompt_text, gen_text), both with special tokens shown.""" + if foundations is None: + foundations = list(_DEFAULT_FORCED_FOUNDATIONS) + schema = _make_forced_hint(foundations) + prompt_text = tok.apply_chat_template( + [{"role": "user", "content": f"{user_prompt}\n\n{schema}"}], + tokenize=False, add_generation_prompt=True) + "\n" + device = next(model.parameters()).device + enc = tok(prompt_text, return_tensors="pt", add_special_tokens=False).to(device) + pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id + gen_kwargs = dict(max_new_tokens=max_think_tokens, pad_token_id=pad_id) + if temperature > 0.0: + gen_kwargs.update(do_sample=True, temperature=temperature, top_p=top_p) + else: + gen_kwargs.update(do_sample=False) + out = model.generate(**enc, **gen_kwargs) + gen_text = tok.decode(out[0, enc.input_ids.shape[1]:], skip_special_tokens=False) + return prompt_text, gen_text +