| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| import os, re, math, time, json, string |
| from collections import Counter |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| DECODING_DEFAULTS = dict( |
| model_name="EleutherAI/pythia-160m", |
| dataset="hamishivi/ROCStories", split="test", n_samples=1000, seed=0, |
| max_new_tokens=70, batch_size=32, beam_batch_size=8, dtype="float32", |
| length_penalty=1.0, |
| configs=[ |
| {"name": "greedy", "strategy": "greedy"}, |
| {"name": "beam_w1", "strategy": "beam", "width": 1}, |
| {"name": "beam_w2", "strategy": "beam", "width": 2}, |
| {"name": "beam_w4", "strategy": "beam", "width": 4}, |
| {"name": "topk_k5", "strategy": "top_k", "k": 5}, |
| {"name": "topk_k20", "strategy": "top_k", "k": 20}, |
| {"name": "topk_k50", "strategy": "top_k", "k": 50}, |
| {"name": "topk_k50_t0.7", "strategy": "top_k", "k": 50, "temperature": 0.7}, |
| {"name": "topp_p0.5", "strategy": "top_p", "p": 0.5}, |
| {"name": "topp_p0.8", "strategy": "top_p", "p": 0.8}, |
| {"name": "topp_p0.95", "strategy": "top_p", "p": 0.95}, |
| {"name": "topp_p0.9_t0.7", "strategy": "top_p", "p": 0.9, "temperature": 0.7}, |
| ], |
| timing_widths=[1, 2, 4, 8], timing_samples=100, |
| judge_model="EleutherAI/pythia-410m", |
| use_bertscore=True, |
| out_dir="/kaggle/working/outputs/decoding" if os.path.isdir("/kaggle/working") else "outputs/decoding", |
| hf_user=None, repo_name="anlp-a2-decoding", push_to_hub=True, private=False, |
| ) |
|
|
|
|
| |
| |
| |
| def load_lm(name, device, dtype="float32"): |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(name) |
| dt = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}[dtype] |
| try: |
| model = AutoModelForCausalLM.from_pretrained(name, dtype=dt) |
| except TypeError: |
| model = AutoModelForCausalLM.from_pretrained(name, torch_dtype=dt) |
| model.to(device).eval() |
| if tok.pad_token_id is None: |
| tok.pad_token = tok.eos_token |
| return model, tok |
|
|
|
|
| def reorder_cache(past, idx): |
| """Select rows (beams) of a KV cache - works for legacy tuples and Cache objects.""" |
| if past is None: |
| return None |
| if isinstance(past, (tuple, list)): |
| return tuple(tuple(t.index_select(0, idx) for t in layer) for layer in past) |
| if hasattr(past, "reorder_cache"): |
| past.reorder_cache(idx) |
| return past |
| if hasattr(past, "layers"): |
| for layer in past.layers: |
| layer.keys = layer.keys.index_select(0, idx) |
| layer.values = layer.values.index_select(0, idx) |
| return past |
| if hasattr(past, "key_cache"): |
| for i in range(len(past.key_cache)): |
| past.key_cache[i] = past.key_cache[i].index_select(0, idx) |
| past.value_cache[i] = past.value_cache[i].index_select(0, idx) |
| return past |
| raise TypeError(f"cannot reorder cache of type {type(past)}") |
|
|
|
|
| @torch.no_grad() |
| def forward_step(model, input_ids, attention_mask, position_ids, past): |
| out = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, |
| past_key_values=past, use_cache=True) |
| return out.logits[:, -1, :].float(), out.past_key_values |
|
|
|
|
| def left_pad(prompts, pad_id, device): |
| L = max(len(p) for p in prompts) |
| ids = torch.full((len(prompts), L), pad_id, dtype=torch.long) |
| mask = torch.zeros((len(prompts), L), dtype=torch.long) |
| for i, p in enumerate(prompts): |
| ids[i, L - len(p):] = torch.tensor(p) |
| mask[i, L - len(p):] = 1 |
| pos = (mask.cumsum(1) - 1).clamp(min=0) |
| return ids.to(device), mask.to(device), pos.to(device) |
|
|
|
|
| |
| |
| |
| def top_k_filter(logits, k): |
| k = min(k, logits.size(-1)) |
| kth = torch.topk(logits, k, dim=-1).values[:, -1:] |
| return logits.masked_fill(logits < kth, float("-inf")) |
|
|
|
|
| def top_p_filter(logits, p): |
| """Keep the smallest set of tokens whose cumulative probability >= p (always >= 1 token).""" |
| sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1) |
| probs = F.softmax(sorted_logits, dim=-1) |
| cum = probs.cumsum(-1) |
| remove = (cum - probs) >= p |
| remove[:, 0] = False |
| sorted_logits = sorted_logits.masked_fill(remove, float("-inf")) |
| return torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits) |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def sample_decode(model, prompts, max_new_tokens, eos_id, pad_id, strategy="greedy", k=None, p=None, |
| temperature=1.0, generator=None): |
| """Returns (list of generated id lists, list of sum log-prob under the *unmodified* |
| model distribution, list of lengths).""" |
| device = next(model.parameters()).device |
| B = len(prompts) |
| ids, mask, pos = left_pad(prompts, pad_id, device) |
| logits, past = forward_step(model, ids, mask, pos, None) |
| next_pos = pos[:, -1:] + 1 |
| finished = torch.zeros(B, dtype=torch.bool, device=device) |
| out, lp_sum, lengths = [], torch.zeros(B, device=device), torch.zeros(B, dtype=torch.long, device=device) |
| for _ in range(max_new_tokens): |
| logp = F.log_softmax(logits, -1) |
| if strategy == "greedy": |
| nxt = logits.argmax(-1) |
| else: |
| lg = logits / max(temperature, 1e-6) |
| if strategy == "top_k": |
| lg = top_k_filter(lg, k) |
| elif strategy == "top_p": |
| lg = top_p_filter(lg, p) |
| nxt = torch.multinomial(F.softmax(lg, -1), 1, generator=generator).squeeze(1) |
| nxt = torch.where(finished, torch.full_like(nxt, pad_id), nxt) |
| lp_sum += torch.where(finished, torch.zeros_like(lp_sum), logp.gather(1, nxt[:, None]).squeeze(1)) |
| lengths += (~finished).long() |
| out.append(nxt) |
| finished |= nxt == eos_id |
| if bool(finished.all()): |
| break |
| mask = torch.cat([mask, torch.ones(B, 1, dtype=mask.dtype, device=device)], 1) |
| logits, past = forward_step(model, nxt[:, None], mask, next_pos, past) |
| next_pos = next_pos + 1 |
| gen = torch.stack(out, 1).tolist() if out else [[] for _ in range(B)] |
| res = [] |
| for row in gen: |
| if eos_id in row: |
| row = row[:row.index(eos_id)] |
| res.append(row) |
| return res, lp_sum.tolist(), lengths.tolist() |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def beam_search(model, prompts, width, max_new_tokens, eos_id, pad_id, length_penalty=1.0): |
| """Standard beam search with a KV cache. |
| * each step: scores(B,W) + log p(B,W,V) -> top 2W candidates per prompt |
| * candidates ending in EOS become finished hypotheses (score / len^alpha) |
| * a prompt is done once it has W finished hypotheses (early stopping) |
| * at max length, open beams are added to the finished pool. |
| Returns (ids, sum log-prob, length) of the best hypothesis per prompt.""" |
| device = next(model.parameters()).device |
| B, W = len(prompts), width |
| rep = [p for p in prompts for _ in range(W)] |
| ids, mask, pos = left_pad(rep, pad_id, device) |
| logits, past = forward_step(model, ids, mask, pos, None) |
| V = logits.size(-1) |
| next_pos = pos[:, -1:] + 1 |
| beam_scores = torch.zeros(B, W, device=device) |
| beam_scores[:, 1:] = float("-inf") |
| seqs = torch.zeros(B * W, 0, dtype=torch.long, device=device) |
| finished = [[] for _ in range(B)] |
| done = [False] * B |
| for t in range(max_new_tokens): |
| logp = F.log_softmax(logits, -1).view(B, W, V) |
| cand = (beam_scores.unsqueeze(-1) + logp).view(B, W * V) |
| top_s, top_i = cand.topk(2 * W, dim=-1) |
| top_s_l, top_i_l = top_s.tolist(), top_i.tolist() |
| new_scores = torch.zeros(B, W, device=device) |
| new_beam = torch.zeros(B, W, dtype=torch.long) |
| new_tok = torch.full((B, W), pad_id, dtype=torch.long) |
| seqs_l = None |
| for b in range(B): |
| if done[b]: |
| new_scores[b] = float("-inf"); new_scores[b, 0] = 0.0 |
| continue |
| j = 0 |
| for r in range(2 * W): |
| s, ci = top_s_l[b][r], top_i_l[b][r] |
| bi, ti = ci // V, ci % V |
| if s == float("-inf"): |
| break |
| if ti == eos_id: |
| if r < W: |
| if seqs_l is None: |
| seqs_l = seqs.tolist() |
| toks = seqs_l[b * W + bi] |
| n = len(toks) + 1 |
| finished[b].append((s / (n ** length_penalty), s, toks, n)) |
| continue |
| new_scores[b, j] = s; new_beam[b, j] = bi; new_tok[b, j] = ti |
| j += 1 |
| if j == W: |
| break |
| while j < W: |
| new_scores[b, j] = float("-inf"); j += 1 |
| if len(finished[b]) >= W: |
| done[b] = True |
| if all(done): |
| break |
| gidx = (torch.arange(B).unsqueeze(1) * W + new_beam).view(-1).to(device) |
| seqs = torch.cat([seqs.index_select(0, gidx), new_tok.view(-1, 1).to(device)], 1) |
| beam_scores = new_scores |
| past = reorder_cache(past, gidx) |
| mask = torch.cat([mask.index_select(0, gidx), torch.ones(B * W, 1, dtype=mask.dtype, device=device)], 1) |
| next_pos = next_pos.index_select(0, gidx) |
| logits, past = forward_step(model, new_tok.view(-1, 1).to(device), mask, next_pos, past) |
| next_pos = next_pos + 1 |
| seqs_l = seqs.tolist(); sc = beam_scores.tolist() |
| res, lps, lens = [], [], [] |
| for b in range(B): |
| if not done[b]: |
| for w in range(W): |
| if sc[b][w] > float("-inf"): |
| toks = seqs_l[b * W + w] |
| n = max(len(toks), 1) |
| finished[b].append((sc[b][w] / (n ** length_penalty), sc[b][w], toks, n)) |
| best = max(finished[b], key=lambda z: z[0]) |
| res.append(best[2]); lps.append(best[1]); lens.append(best[3]) |
| return res, lps, lens |
|
|
|
|
| |
| |
| |
| _PUNCT = re.compile(f"[{re.escape(string.punctuation)}]") |
|
|
|
|
| def norm_words(s): |
| return _PUNCT.sub(" ", s.lower()).split() |
|
|
|
|
| def clean_generation(text, max_sentences=4): |
| """Post-processing applied before scoring: cut at the first newline (Pythia tends to |
| drift into unrelated text after a line break) and keep at most 4 sentences |
| (ROCStories continuations are exactly 4 sentences).""" |
| text = text.lstrip() |
| text = text.split("\n")[0] |
| ends = [m.end() for m in re.finditer(r"[.!?](?=\s|$)", text)] |
| if len(ends) >= max_sentences: |
| text = text[:ends[max_sentences - 1]] |
| return " " + text.strip() |
|
|
|
|
| def token_prf_acc(gen, ref): |
| g, r = norm_words(gen), norm_words(ref) |
| common = Counter(g) & Counter(r) |
| ov = sum(common.values()) |
| P = ov / len(g) if g else 0.0 |
| R = ov / len(r) if r else 0.0 |
| F1 = 2 * P * R / (P + R) if P + R > 0 else 0.0 |
| acc = sum(1 for i in range(len(r)) if i < len(g) and g[i] == r[i]) / len(r) if r else 0.0 |
| return P, R, F1, acc |
|
|
|
|
| def distinct_n(texts, n): |
| grams, total = set(), 0 |
| for t in texts: |
| w = norm_words(t) |
| ng = [tuple(w[i:i + n]) for i in range(len(w) - n + 1)] |
| grams.update(ng); total += len(ng) |
| return len(grams) / max(total, 1) |
|
|
|
|
| def repetition_rate(text, n=4): |
| w = norm_words(text) |
| ng = [tuple(w[i:i + n]) for i in range(len(w) - n + 1)] |
| return 1 - len(set(ng)) / len(ng) if ng else 0.0 |
|
|
|
|
| @torch.no_grad() |
| def continuation_ppl(model, prompt_ids, cont_ids, pad_id, batch_size=32): |
| """Per-sample perplexity of cont given prompt (teacher forced).""" |
| device = next(model.parameters()).device |
| out = [] |
| for s in range(0, len(prompt_ids), batch_size): |
| P = prompt_ids[s:s + batch_size]; C = cont_ids[s:s + batch_size] |
| seqs = [p + c for p, c in zip(P, C)] |
| L = max(len(x) for x in seqs) |
| ids = torch.full((len(seqs), L), pad_id, dtype=torch.long) |
| att = torch.zeros((len(seqs), L), dtype=torch.long) |
| lab = torch.full((len(seqs), L), -100, dtype=torch.long) |
| for i, (p, c) in enumerate(zip(P, C)): |
| x = p + c |
| ids[i, :len(x)] = torch.tensor(x); att[i, :len(x)] = 1 |
| lab[i, len(p):len(x)] = torch.tensor(c) if c else lab[i, len(p):len(x)] |
| ids, att, lab = ids.to(device), att.to(device), lab.to(device) |
| logits = model(input_ids=ids, attention_mask=att).logits.float() |
| lp = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)), lab[:, 1:].reshape(-1), |
| ignore_index=-100, reduction="none").view(len(seqs), -1) |
| n = (lab[:, 1:] != -100).sum(1).clamp(min=1) |
| out += torch.exp(lp.sum(1) / n).tolist() |
| return out |
|
|
|
|
| def score_config(gens, refs, self_lp, self_len): |
| import sacrebleu |
| prf = np.array([token_prf_acc(g, r) for g, r in zip(gens, refs)]) |
| ppl_self = [math.exp(-lp / max(n, 1)) for lp, n in zip(self_lp, self_len)] |
| return { |
| "precision": float(prf[:, 0].mean()), "recall": float(prf[:, 1].mean()), |
| "f1": float(prf[:, 2].mean()), "accuracy": float(prf[:, 3].mean()), |
| "bleu": sacrebleu.corpus_bleu(gens, [refs]).score, |
| "self_ppl_mean": float(np.mean(ppl_self)), "self_ppl_median": float(np.median(ppl_self)), |
| "distinct_1": distinct_n(gens, 1), "distinct_2": distinct_n(gens, 2), |
| "rep_4gram": float(np.mean([repetition_rate(g) for g in gens])), |
| "avg_words": float(np.mean([len(norm_words(g)) for g in gens])), |
| } |
|
|
|
|
| |
| |
| |
| def run_config(model, tok, prompts, conf, cfg, seed=0): |
| device = next(model.parameters()).device |
| eos, pad = tok.eos_token_id, tok.pad_token_id |
| gen_ids, lps, lens = [None] * len(prompts), [None] * len(prompts), [None] * len(prompts) |
| order = sorted(range(len(prompts)), key=lambda i: len(prompts[i])) |
| g = torch.Generator(device=device); g.manual_seed(seed) |
| bs = cfg["beam_batch_size"] if conf["strategy"] == "beam" else cfg["batch_size"] |
| if device.type == "cuda": torch.cuda.synchronize() |
| t0 = time.time() |
| for s in range(0, len(order), bs): |
| idx = order[s:s + bs] |
| P = [prompts[i] for i in idx] |
| if conf["strategy"] == "beam": |
| r, lp, ln = beam_search(model, P, conf["width"], cfg["max_new_tokens"], eos, pad, cfg["length_penalty"]) |
| else: |
| r, lp, ln = sample_decode(model, P, cfg["max_new_tokens"], eos, pad, conf["strategy"], |
| conf.get("k"), conf.get("p"), conf.get("temperature", 1.0), g) |
| for i, a, b, c in zip(idx, r, lp, ln): |
| gen_ids[i], lps[i], lens[i] = a, b, c |
| if device.type == "cuda": torch.cuda.synchronize() |
| return gen_ids, lps, lens, time.time() - t0 |
|
|
|
|
| def beam_timing(model, tok, prompts, widths, max_new_tokens): |
| """Wall time of beam search with batch size 1 (one prompt at a time) for each width. |
| Two modes: 'natural' (stops early once W hypotheses hit EOS) and 'fixed' (EOS disabled, |
| always max_new_tokens steps) - the fixed mode isolates the per-step cost as a function of W.""" |
| device = next(model.parameters()).device |
| res = [] |
| beam_search(model, prompts[:2], 2, 5, tok.eos_token_id, tok.pad_token_id) |
| for mode, eos in (("natural", tok.eos_token_id), ("fixed", -1)): |
| for w in widths: |
| times, steps = [], [] |
| for p in prompts: |
| if device.type == "cuda": torch.cuda.synchronize() |
| t0 = time.time() |
| r, _, ln = beam_search(model, [p], w, max_new_tokens, eos, tok.pad_token_id) |
| if device.type == "cuda": torch.cuda.synchronize() |
| times.append(time.time() - t0); steps.append(max(len(r[0]), 1)) |
| rec = {"mode": mode, "width": w, "mean_s_per_prompt": float(np.mean(times)), "std_s": float(np.std(times)), |
| "mean_ms_per_token": float(np.mean([t / s for t, s in zip(times, steps)]) * 1e3), |
| "mean_output_tokens": float(np.mean(steps)), "total_s": float(np.sum(times)), |
| "n_prompts": len(prompts)} |
| res.append(rec) |
| print(f" [{mode}] width {w}: {rec['mean_s_per_prompt']*1e3:.1f} ms/prompt, " |
| f"{rec['mean_ms_per_token']:.2f} ms/token, {rec['mean_output_tokens']:.1f} tokens") |
| return res |
|
|
|
|
| def run_decoding_eval(user_cfg=None): |
| import sys |
| sys.path.insert(0, ".") |
| from common import get_hf_token, hf_login, hf_resolve_user, hf_push_folder, save_json |
| cfg = dict(DECODING_DEFAULTS); cfg.update(user_cfg or {}) |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| os.makedirs(cfg["out_dir"], exist_ok=True) |
| torch.manual_seed(cfg["seed"]) |
|
|
| token = get_hf_token(); hf_login(token) |
| model, tok = load_lm(cfg["model_name"], device, cfg["dtype"]) |
| print(f"loaded {cfg['model_name']} on {device}; eos={tok.eos_token_id} pad={tok.pad_token_id}") |
|
|
| if cfg.get("dataset_override") is not None: |
| rows = cfg["dataset_override"] |
| else: |
| from datasets import load_dataset |
| rows = load_dataset(cfg["dataset"], split=cfg["split"]) |
| prompts_txt = list(rows["prompt"]); refs_txt = list(rows["continuation"]) |
| keep = [i for i, r in enumerate(refs_txt) if r and r.strip()] |
| rng = np.random.default_rng(cfg["seed"]) |
| keep = sorted(rng.choice(keep, min(cfg["n_samples"], len(keep)), replace=False).tolist()) |
| prompts_txt = [prompts_txt[i] for i in keep]; refs_txt = [refs_txt[i] for i in keep] |
| refs_txt = [" " + r.strip() for r in refs_txt] |
| prompts = [tok(p)["input_ids"] for p in prompts_txt] |
| print(f"{len(prompts)} test prompts") |
|
|
| |
| ref_ids = [tok(r)["input_ids"] for r in refs_txt] |
| ref_ppl = continuation_ppl(model, prompts, ref_ids, tok.pad_token_id) |
| summary = {"reference_ppl_mean": float(np.mean(ref_ppl)), "reference_ppl_median": float(np.median(ref_ppl)), |
| "n_samples": len(prompts), "model": cfg["model_name"], "max_new_tokens": cfg["max_new_tokens"]} |
| print(f"reference continuation ppl: mean {summary['reference_ppl_mean']:.2f}, " |
| f"median {summary['reference_ppl_median']:.2f}") |
|
|
| results, generations = [], {"prompt": prompts_txt, "reference": refs_txt} |
| for conf in cfg["configs"]: |
| print(f"== {conf['name']} ==") |
| ids, lps, lens, dt = run_config(model, tok, prompts, conf, cfg, cfg["seed"]) |
| raw = [tok.decode(x, skip_special_tokens=True) for x in ids] |
| gens = [clean_generation(t) for t in raw] |
| m = score_config(gens, refs_txt, lps, lens) |
| m.update({"name": conf["name"], "strategy": conf["strategy"], "params": conf, |
| "time_s": dt, "ms_per_prompt": dt / len(prompts) * 1e3, |
| "gen_tokens_mean": float(np.mean([len(x) for x in ids]))}) |
| results.append(m) |
| generations[conf["name"]] = gens |
| generations[conf["name"] + "__raw"] = raw |
| print(f" P {m['precision']:.3f} R {m['recall']:.3f} F1 {m['f1']:.3f} Acc {m['accuracy']:.3f} " |
| f"BLEU {m['bleu']:.2f} self-ppl {m['self_ppl_median']:.2f} dist-2 {m['distinct_2']:.3f} " |
| f"rep4 {m['rep_4gram']:.3f} time {dt:.1f}s") |
|
|
| |
| if "greedy" in generations and "beam_w1" in generations: |
| same = np.mean([a == b for a, b in zip(generations["greedy__raw"], generations["beam_w1__raw"])]) |
| summary["beam_w1_equals_greedy_frac"] = float(same) |
| print(f"beam(w=1) identical to greedy on {same*100:.1f}% of prompts") |
|
|
| |
| if cfg["judge_model"]: |
| try: |
| judge, _ = load_lm(cfg["judge_model"], device, cfg["dtype"]) |
| for m in results: |
| gen_ids = [tok(g)["input_ids"] for g in generations[m["name"]]] |
| pp = continuation_ppl(judge, prompts, gen_ids, tok.pad_token_id) |
| m["judge_ppl_median"] = float(np.median(pp)); m["judge_ppl_mean"] = float(np.mean(pp)) |
| jr = continuation_ppl(judge, prompts, ref_ids, tok.pad_token_id) |
| summary["reference_judge_ppl_median"] = float(np.median(jr)) |
| del judge |
| except Exception as e: |
| print("[judge] skipped:", e) |
|
|
| if cfg["use_bertscore"]: |
| try: |
| from bert_score import score as bscore |
| for m in results: |
| P, R, F1 = bscore(generations[m["name"]], refs_txt, lang="en", verbose=False, |
| device=str(device), batch_size=64) |
| m["bertscore_p"], m["bertscore_r"], m["bertscore_f1"] = (float(P.mean()), float(R.mean()), |
| float(F1.mean())) |
| except Exception as e: |
| print("[bertscore] skipped:", e) |
|
|
| print("beam-search timing (batch size 1):") |
| timing = beam_timing(model, tok, prompts[:cfg["timing_samples"]], cfg["timing_widths"], cfg["max_new_tokens"]) |
|
|
| save_json({"summary": summary, "results": results, "beam_timing": timing, "config": |
| {k: v for k, v in cfg.items() if k != "dataset_override"}}, |
| os.path.join(cfg["out_dir"], "decoding_results.json")) |
| save_json(generations, os.path.join(cfg["out_dir"], "generations.json")) |
| shutil_copy_code(cfg["out_dir"]) |
| user = hf_resolve_user(cfg["hf_user"], token) |
| if cfg["push_to_hub"] and user: |
| hf_push_folder(cfg["out_dir"], f"{user}/{cfg['repo_name']}", token, cfg["private"], "decoding results") |
| print(f"pushed to https://huggingface.co/{user}/{cfg['repo_name']}") |
| return summary, results, timing, generations |
|
|
|
|
| def shutil_copy_code(out_dir): |
| import shutil |
| for f in ("decoding.py", "common.py"): |
| if os.path.exists(f): |
| shutil.copy(f, os.path.join(out_dir, f)) |
|
|
|
|
| |
| |
| |
| PALETTE = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#4a3aa7", "#e34948"] |
| STRATEGY_COLOR = {"greedy": PALETTE[0], "beam": PALETTE[6], "top_k": PALETTE[1], "top_p": PALETTE[2]} |
|
|
|
|
| def decoding_report(res_json_path, out_dir): |
| import pandas as pd |
| import matplotlib.pyplot as plt |
| from common import setup_plot_style |
| setup_plot_style() |
| R = json.load(open(res_json_path)) |
| df = pd.DataFrame(R["results"]) |
| cols = ["name", "accuracy", "precision", "recall", "f1", "bleu", "self_ppl_median", "judge_ppl_median", |
| "bertscore_f1", "distinct_2", "rep_4gram", "avg_words", "ms_per_prompt"] |
| table = df[[c for c in cols if c in df.columns]].copy() |
| table.to_csv(os.path.join(out_dir, "decoding_table.csv"), index=False) |
| colors = [STRATEGY_COLOR[s] for s in df["strategy"]] |
|
|
| metrics = [("f1", "token F1"), ("accuracy", "positional accuracy"), ("bleu", "BLEU"), |
| ("self_ppl_median", "self perplexity (median)"), ("distinct_2", "distinct-2"), |
| ("rep_4gram", "repeated 4-gram rate")] |
| fig, axes = plt.subplots(2, 3, figsize=(15, 7.5), constrained_layout=True) |
| for ax, (k, lab) in zip(axes.flat, metrics): |
| ax.bar(range(len(df)), df[k], color=colors, width=0.7) |
| ax.set_xticks(range(len(df))); ax.set_xticklabels(df["name"], rotation=60, ha="right", fontsize=8) |
| ax.set_title(lab); ax.grid(axis="x", visible=False) |
| handles = [plt.Rectangle((0, 0), 1, 1, color=c) for c in STRATEGY_COLOR.values()] |
| fig.legend(handles, STRATEGY_COLOR.keys(), loc="upper right", ncol=4) |
| fig.savefig(os.path.join(out_dir, "decoding_metrics.png")) |
|
|
| |
| fig2, ax = plt.subplots(figsize=(6.5, 4.5), constrained_layout=True) |
| for _, r in df.iterrows(): |
| ax.scatter(r["distinct_2"], r["f1"], color=STRATEGY_COLOR[r["strategy"]], s=60, zorder=3, |
| edgecolor="white", linewidth=1.5) |
| ax.annotate(r["name"], (r["distinct_2"], r["f1"]), fontsize=7, xytext=(4, 3), textcoords="offset points") |
| ax.set_xlabel("distinct-2 (diversity)"); ax.set_ylabel("token F1 vs reference") |
| ax.set_title("quality / diversity trade-off") |
| fig2.savefig(os.path.join(out_dir, "quality_vs_diversity.png")) |
|
|
| |
| tm = pd.DataFrame(R["beam_timing"]) |
| fig3, axes3 = plt.subplots(1, 2, figsize=(11, 3.8), constrained_layout=True) |
| for i, (mode, g) in enumerate(tm.groupby("mode", sort=False)): |
| axes3[0].plot(g["width"], g["mean_s_per_prompt"] * 1e3, marker="o", label=mode, color=PALETTE[i]) |
| axes3[1].plot(g["width"], g["mean_ms_per_token"], marker="o", label=mode, color=PALETTE[i]) |
| for ax, lab in zip(axes3, ["ms per prompt", "ms per generated token"]): |
| ax.set_xlabel("beam width"); ax.set_ylabel(lab); ax.set_xticks(sorted(tm["width"].unique())); ax.legend() |
| axes3[0].set_title("beam search wall time (batch size 1)") |
| axes3[1].set_title("per-step cost") |
| fig3.savefig(os.path.join(out_dir, "beam_timing.png")) |
| return table, tm |
|
|
|
|
| def show_examples(gen_json_path, n=5, names=None, seed=0): |
| G = json.load(open(gen_json_path)) |
| names = names or [k for k in G if k not in ("prompt", "reference") and not k.endswith("__raw")] |
| rng = np.random.default_rng(seed) |
| for i in rng.choice(len(G["prompt"]), n, replace=False): |
| print("=" * 100) |
| print("PROMPT :", G["prompt"][i]) |
| print("REFERENCE:", G["reference"][i].strip()) |
| for k in names: |
| print(f"{k:>16s}: {G[k][i].strip()}") |
|
|
