| from __future__ import annotations
|
|
|
| import argparse
|
| from collections import Counter
|
| import json
|
| import math
|
| from pathlib import Path
|
| from typing import Any
|
| import torch
|
| import torch.nn.functional as F
|
| from torch.utils.data import DataLoader
|
|
|
| from conveyer.checkpoint import load_checkpoint
|
| from conveyer.config import ConveyerConfig
|
| from conveyer.data import build_dataset_and_tokenizer, collate_conveyer_batch
|
| from conveyer.infer import generate_conveyer
|
| from conveyer.losses import compute_unrolled_conveyer_loss
|
| from conveyer.model import ConveyerDiffusionLM
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| COMMONSENSE_CLOZE_PAIRS: list[tuple[str, str, str, str]] = [
|
| (
|
| "emotion_cause",
|
| "Once upon a time, Tom lost his favorite red ball in the dark woods. Because he lost his ball, Tom felt very",
|
| " sad and started to cry.",
|
| " happy and started to laugh.",
|
| ),
|
| (
|
| "emotion_cause",
|
| "Lily worked all morning to bake a sweet cake for her mom. When her mom smiled and hugged her, Lily felt very",
|
| " proud and happy.",
|
| " angry and mean.",
|
| ),
|
| (
|
| "physical_weight",
|
| "The little mouse tried to lift the huge heavy rock, but the rock was much too",
|
| " heavy for the tiny mouse to move.",
|
| " light and floated into the sky.",
|
| ),
|
| (
|
| "physical_temperature",
|
| "Ben went outside in the snow without his coat or gloves. Soon his hands felt very",
|
| " cold and freezing.",
|
| " hot and burning.",
|
| ),
|
| (
|
| "physical_temperature",
|
| "The soup had just come off the stove and was boiling. Mom told Mia to wait because the soup was too",
|
| " hot to eat right now.",
|
| " cold and frozen like ice.",
|
| ),
|
| (
|
| "physical_fragile",
|
| "The glass cup fell off the high table onto the hard stone floor. With a loud crash, the glass",
|
| " broke into many small pieces.",
|
| " bounced high like a rubber ball.",
|
| ),
|
| (
|
| "object_use",
|
| "The wooden door was locked tight. To unlock the door, Anna took out her small brass",
|
| " key and turned it in the lock.",
|
| " spoon and stirred the soup.",
|
| ),
|
| (
|
| "object_use",
|
| "It started to rain hard outside. Before walking to school, Max opened his big",
|
| " umbrella to stay dry.",
|
| " pillow to go to sleep.",
|
| ),
|
| (
|
| "biological_need",
|
| "The little puppy had been running in the park all day and had not eaten anything. His tummy rumbled because he was very",
|
| " hungry and wanted some food.",
|
| " sleepy and wanted a haircut.",
|
| ),
|
| (
|
| "biological_need",
|
| "After playing all day until night came and the stars shone in the sky, the tired children went to bed to",
|
| " sleep until morning.",
|
| " swim in the cold ocean.",
|
| ),
|
| (
|
| "entity_tracking",
|
| "Sam had a blue car and Lucy had a yellow doll. Sam gave his blue car to Lucy. Now Lucy had the yellow doll and the",
|
| " blue car.",
|
| " green bicycle.",
|
| ),
|
| (
|
| "entity_tracking",
|
| " Lily put her shiny coin inside the small red box and closed the lid. When she opened the red box later, she saw her",
|
| " shiny coin inside.",
|
| " big brown horse inside.",
|
| ),
|
| (
|
| "danger_reaction",
|
| "Lily and a little bunny walked into the green forest. Suddenly a big fierce lion jumped out and roared loudly. The little bunny was",
|
| " scared and hid behind a tree.",
|
| " bored and read a newspaper.",
|
| ),
|
| (
|
| "social_repair",
|
| "Tom accidentally knocked over Mia's tall block tower. Tom felt bad about the accident, so he said,",
|
| ' "I am sorry, Mia. Let me help you build it again."',
|
| ' "Thank you for the delicious birthday cake."',
|
| ),
|
| (
|
| "spatial_logic",
|
| "The bird wanted to reach the nest at the very top of the tall tree, so it flapped its wings and",
|
| " flew up high into the branches.",
|
| " dug a deep hole under the mud.",
|
| ),
|
| (
|
| "weather_visibility",
|
| "At midnight all the lights went out in the room. Because it was pitch dark, Ben could not",
|
| " see anything around him.",
|
| " hear the loud thunder outside.",
|
| ),
|
| ]
|
|
|
|
|
| NOVEL_STORY_PROMPTS: list[tuple[str, list[str]]] = [
|
| (
|
| "Once upon a time, Lily and a little bunny walked into the green forest to find the golden key near the locked treasure chest. Suddenly a big lion",
|
| ["lion", "Lily", "bunny", "key", "chest", "forest", "scared", "roar", "said"],
|
| ),
|
| (
|
| "Once upon a time, there was a little robot named Beep who lived in a noisy workshop. One rainy morning, Beep found a tiny lost kitten by the door. Beep",
|
| ["Beep", "kitten", "robot", "door", "rain", "warm", "milk", "help"],
|
| ),
|
| (
|
| "Sam and his sister Mia built a huge sandcastle on the beach near the waves. Suddenly, a big wave splashed toward their castle. Sam quickly",
|
| ["Sam", "Mia", "sand", "castle", "wave", "beach", "water", "wall"],
|
| ),
|
| (
|
| "In a quiet village, an old baker dropped his round loaf of bread and it rolled out the door down the hill. A hungry puppy saw the bread and",
|
| ["bread", "puppy", "baker", "hill", "chased", "ran", "ate", "mouth"],
|
| ),
|
| (
|
| "Lucy had a box of bright crayons, but she could not find the blue crayon to color the sky. She looked under her bed and",
|
| ["Lucy", "crayon", "blue", "bed", "sky", "found", "draw", "paper"],
|
| ),
|
| (
|
| "One winter day, Max made a snowman in the yard with a carrot nose. When the warm sun came out in the afternoon, the snowman began to",
|
| ["Max", "snowman", "sun", "melt", "carrot", "snow", "warm", "water"],
|
| ),
|
| ]
|
|
|
|
|
| @torch.no_grad()
|
| def evaluate_zero_shot_commonsense(
|
| model: ConveyerDiffusionLM,
|
| tokenizer: Any,
|
| device: torch.device,
|
| ) -> dict[str, Any]:
|
| """
|
| Evaluates zero-shot 2-way multiple-choice accuracy (BLiMP / StoryCloze style) using both:
|
| 1. Causal LM log-probability per token over the continuation.
|
| 2. Conveyor Slot-0 first-token preference on the continuation.
|
| """
|
| model.eval()
|
| correct_lm = 0
|
| correct_s0 = 0
|
| margin_sum = 0.0
|
| details: list[dict[str, Any]] = []
|
|
|
| for category, ctx, good_cont, bad_cont in COMMONSENSE_CLOZE_PAIRS:
|
| ctx_ids = tokenizer.encode(ctx.strip())
|
| good_ids = tokenizer.encode(ctx.strip() + good_cont)
|
| bad_ids = tokenizer.encode(ctx.strip() + bad_cont)
|
| plen = len(ctx_ids)
|
|
|
| def _score_seq(full_ids: list[int]) -> float:
|
| inp = torch.tensor([full_ids], dtype=torch.long, device=device)
|
| logits, _ = model.forward_causal_and_cache(inp, prefix_len=len(full_ids))
|
| log_probs = F.log_softmax(logits[0, :-1, :].float(), dim=-1)
|
| targets = inp[0, 1:]
|
|
|
| start_idx = max(0, plen - 1)
|
| cont_lp = log_probs[start_idx:, :].gather(1, targets[start_idx:].unsqueeze(1)).squeeze(1)
|
| return float(cont_lp.mean().item())
|
|
|
| score_good = _score_seq(good_ids)
|
| score_bad = _score_seq(bad_ids)
|
| is_ok = score_good > score_bad
|
| correct_lm += int(is_ok)
|
| margin_sum += score_good - score_bad
|
|
|
|
|
| if len(good_ids) > plen and len(bad_ids) > plen:
|
| p_tensor = torch.tensor([ctx_ids], dtype=torch.long, device=device)
|
| kv = model.prefill_prefix(p_tensor)
|
| belt = model.init_empty_belt(1, tick=plen, device=device, deterministic=True)
|
| out = model.forward_belt(belt, kv_cache=kv)
|
| s0_lp = F.log_softmax(out.slot0_logits[0].float(), dim=-1)
|
| s0_good = float(s0_lp[good_ids[plen]].item())
|
| s0_bad = float(s0_lp[bad_ids[plen]].item())
|
| correct_s0 += int(s0_good > s0_bad)
|
|
|
| details.append(
|
| {
|
| "category": category,
|
| "correct": is_ok,
|
| "logprob_good": round(score_good, 3),
|
| "logprob_bad": round(score_bad, 3),
|
| }
|
| )
|
|
|
| n = max(1, len(COMMONSENSE_CLOZE_PAIRS))
|
| return {
|
| "commonsense_lm_acc": correct_lm / n,
|
| "commonsense_slot0_acc": correct_s0 / n,
|
| "mean_logprob_margin": margin_sum / n,
|
| "num_pairs": len(COMMONSENSE_CLOZE_PAIRS),
|
| "details": details,
|
| }
|
|
|
|
|
| @torch.no_grad()
|
| def evaluate_prefix_sensitivity_and_generation_quality(
|
| model: ConveyerDiffusionLM,
|
| tokenizer: Any,
|
| device: torch.device,
|
| max_new_tokens: int = 80,
|
| ) -> dict[str, Any]:
|
| """
|
| Measures:
|
| 1. Prefix Sensitivity KL(P(slot0 | prompt) || P(slot0 | shuffled_prompt)):
|
| In the broken template model this was 0.03 nats (ignored word order).
|
| A real language model should have KL >> 0.5 nats.
|
| 2. Open-Ended Generation Quality on 6 unseen prompts:
|
| - Distinct-2 / Distinct-3 n-gram diversity (1.0 = no repetition, <0.5 = looping)
|
| - 4-gram repetition rate (0.0 = clean prose)
|
| - Prompt entity/topic retention rate (% of continuations that reference the prompt's new entities/events)
|
| """
|
| model.eval()
|
| kl_list: list[float] = []
|
| dist2_list: list[float] = []
|
| dist3_list: list[float] = []
|
| rep4_list: list[float] = []
|
| entity_hits = 0
|
| samples: list[dict[str, str]] = []
|
|
|
| g = torch.Generator().manual_seed(123)
|
| for prompt_text, expected_keywords in NOVEL_STORY_PROMPTS:
|
| ids = tokenizer.encode(prompt_text.strip())
|
| p_true = torch.tensor([ids], dtype=torch.long, device=device)
|
| kv_true = model.prefill_prefix(p_true)
|
| belt_true = model.init_empty_belt(1, tick=len(ids), device=device, deterministic=True)
|
| d_true = F.softmax(model.forward_belt(belt_true, kv_cache=kv_true).slot0_logits[0].float(), dim=-1)
|
|
|
| perm = torch.randperm(len(ids) - 1, generator=g) + 1
|
| shuf_ids = [ids[0]] + [ids[i] for i in perm.tolist()]
|
| p_shuf = torch.tensor([shuf_ids], dtype=torch.long, device=device)
|
| kv_shuf = model.prefill_prefix(p_shuf)
|
| belt_shuf = model.init_empty_belt(1, tick=len(ids), device=device, deterministic=True)
|
| d_shuf = F.softmax(model.forward_belt(belt_shuf, kv_cache=kv_shuf).slot0_logits[0].float(), dim=-1)
|
|
|
| kl = float((d_true * (d_true.clamp_min(1e-9).log() - d_shuf.clamp_min(1e-9).log())).sum().item())
|
| kl_list.append(kl)
|
|
|
| out_ids, _ = generate_conveyer(
|
| model=model,
|
| prompt_tokens=p_true,
|
| max_new_tokens=max_new_tokens,
|
| prompt_overlap=10,
|
| temperature=0.0,
|
| eos_token_id=tokenizer.eos_id,
|
| deterministic_spawn=True,
|
| )
|
| gen_toks = out_ids[0, len(ids) :].tolist()
|
| gen_text = tokenizer.decode(gen_toks).strip()
|
|
|
| words = [w.strip(".,!?;:\"'()").lower() for w in gen_text.split() if w.strip()]
|
| if len(words) >= 4:
|
| bigrams = list(zip(words[:-1], words[1:]))
|
| trigrams = list(zip(words[:-2], words[1:-1], words[2:]))
|
| fourgrams = list(zip(words[:-3], words[1:-2], words[2:-1], words[3:]))
|
| dist2_list.append(len(set(bigrams)) / max(1, len(bigrams)))
|
| dist3_list.append(len(set(trigrams)) / max(1, len(trigrams)))
|
| c4 = Counter(fourgrams)
|
| repeated_4 = sum(v - 1 for v in c4.values() if v > 1)
|
| rep4_list.append(repeated_4 / max(1, len(fourgrams)))
|
|
|
| low_gen = gen_text.lower()
|
| hit = any(kw.lower() in low_gen for kw in expected_keywords[:4])
|
| entity_hits += int(hit)
|
| samples.append({"prompt": prompt_text, "continuation": gen_text})
|
|
|
| return {
|
| "prefix_word_order_kl": sum(kl_list) / max(1, len(kl_list)),
|
| "distinct_2": sum(dist2_list) / max(1, len(dist2_list)),
|
| "distinct_3": sum(dist3_list) / max(1, len(dist3_list)),
|
| "repeat_4gram_rate": sum(rep4_list) / max(1, len(rep4_list)),
|
| "prompt_entity_retention": entity_hits / max(1, len(NOVEL_STORY_PROMPTS)),
|
| "samples": samples,
|
| }
|
|
|
|
|
| @torch.no_grad()
|
| def evaluate_crystallization_curve(
|
| model: ConveyerDiffusionLM,
|
| cfg: ConveyerConfig,
|
| val_loader: DataLoader[dict[str, torch.Tensor]],
|
| device: torch.device,
|
| pad_id: int = 0,
|
| max_batches: int = 20,
|
| ) -> dict[str, Any]:
|
| """
|
| Computes the Per-Slot Accuracy, Entropy, and Pause Rate across the 64-slot oven belt
|
| plus held-out real TinyStories Perplexity (Causal LM & Conveyor Slot-0).
|
| """
|
| model.eval()
|
| B = model.belt_size
|
| acc_sum = torch.zeros(B, device=device)
|
| ent_sum = torch.zeros(B, device=device)
|
| s0_loss_sum = 0.0
|
| lm_loss_sum = 0.0
|
| lm_acc_sum = 0.0
|
| pause_sum = 0.0
|
| count = 0
|
|
|
| for i, batch in enumerate(val_loader):
|
| if i >= max_batches:
|
| break
|
| seq = batch["input_ids"].to(device)
|
| plen = batch["prompt_lengths"].to(device)
|
| use_amp = device.type == "cuda" and torch.cuda.is_bf16_supported()
|
| with torch.autocast(device_type=device.type, dtype=torch.bfloat16 if use_amp else torch.float32, enabled=use_amp):
|
| m = compute_unrolled_conveyer_loss(
|
| model=model,
|
| seq_tokens=seq,
|
| cfg=cfg,
|
| prompt_lengths=plen,
|
| tick_start=int(plen.min().item()),
|
| pad_token_id=pad_id,
|
| )
|
| acc_sum += m.per_slot_acc
|
| ent_sum += m.per_slot_entropy
|
| s0_loss_sum += m.slot0_loss
|
| lm_loss_sum += m.prefix_lm_loss
|
| lm_acc_sum += m.prefix_lm_acc
|
| pause_sum += m.mean_plan_pause
|
| count += 1
|
|
|
| denom = max(count, 1)
|
| acc_curve = (acc_sum / denom).cpu().tolist()
|
| ent_curve = (ent_sum / denom).cpu().tolist()
|
| s0_ce = s0_loss_sum / denom
|
| lm_ce = lm_loss_sum / denom
|
| return {
|
| "slot0_ce": s0_ce,
|
| "slot0_ppl": math.exp(min(s0_ce, 20.0)),
|
| "slot0_acc": acc_curve[0],
|
| "causal_lm_ce": lm_ce,
|
| "causal_lm_ppl": math.exp(min(lm_ce, 20.0)),
|
| "causal_lm_acc": lm_acc_sum / denom,
|
| "mean_plan_pause": pause_sum / denom,
|
| "per_slot_acc": acc_curve,
|
| "per_slot_entropy": ent_curve,
|
| }
|
|
|
|
|
| @torch.no_grad()
|
| def evaluate_closed_loop_generation(
|
| model: ConveyerDiffusionLM,
|
| val_ds: Any,
|
| tokenizer: Any,
|
| device: torch.device,
|
| num_episodes: int = 32,
|
| prompt_overlap: int = 12,
|
| bake_steps: int = 64,
|
| max_new_tokens: int = 24,
|
| ) -> dict[str, float]:
|
| model.eval()
|
| n_eval = min(num_episodes, len(val_ds))
|
| first_tok_correct = 0
|
| total_tok_correct = 0
|
| total_tok_count = 0
|
| exact_match = 0
|
|
|
| for i in range(n_eval):
|
| item = val_ds[i]
|
| seq = item["input_ids"].to(device)
|
| plen = int(item["prompt_len"])
|
| prompt = seq[:plen].unsqueeze(0)
|
|
|
| gt_full = seq[plen:].tolist()
|
| gt_target: list[int] = []
|
| for tid in gt_full:
|
| if tid == tokenizer.pad_id:
|
| break
|
| gt_target.append(tid)
|
| if tid == tokenizer.eos_id:
|
| break
|
|
|
| if not gt_target:
|
| continue
|
|
|
| gen_len = min(max_new_tokens, len(gt_target))
|
| full_out, _ = generate_conveyer(
|
| model=model,
|
| prompt_tokens=prompt,
|
| max_new_tokens=gen_len,
|
| prompt_overlap=prompt_overlap,
|
| bake_steps=bake_steps,
|
| temperature=0.0,
|
| eos_token_id=tokenizer.eos_id,
|
| deterministic_spawn=True,
|
| )
|
| pred_ids = full_out[0, plen : plen + gen_len].tolist()
|
| gt_slice = gt_target[:gen_len]
|
|
|
| if len(pred_ids) > 0 and pred_ids[0] == gt_slice[0]:
|
| first_tok_correct += 1
|
|
|
| matches = sum(1 for p, g in zip(pred_ids, gt_slice) if p == g)
|
| total_tok_correct += matches
|
| total_tok_count += len(gt_slice)
|
| if pred_ids == gt_slice:
|
| exact_match += 1
|
|
|
| return {
|
| "first_token_acc": first_tok_correct / max(1, n_eval),
|
| "token_acc": total_tok_correct / max(1, total_tok_count),
|
| "exact_seq_acc": exact_match / max(1, n_eval),
|
| }
|
|
|
|
|
| def main() -> None:
|
| parser = argparse.ArgumentParser(description="Comprehensive Benchmark Suite for Conveyor-Belt Diffusion LLM")
|
| parser.add_argument("--ckpt", type=str, default="runs/tinystories/best.pt", help="Path to checkpoint (.pt)")
|
| parser.add_argument("--device", type=str, default="auto", help="'cpu', 'cuda', or 'auto'")
|
| parser.add_argument("--episodes", type=int, default=32, help="Number of closed-loop episodes")
|
| parser.add_argument("--bake-steps", type=int, nargs="+", default=[16, 64], help="Bake depths K to evaluate")
|
| parser.add_argument("--overlaps", type=int, nargs="+", default=[0, 10], help="Prompt-overlap windows W_pre")
|
| parser.add_argument("--out-json", type=str, default=None, help="Optional path to save evaluation JSON")
|
| args = parser.parse_args()
|
|
|
| if args.device == "auto":
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| else:
|
| device = torch.device(args.device)
|
|
|
| model, cfg, ckpt_meta = load_checkpoint(args.ckpt, device=device)
|
| _, val_ds, tokenizer = build_dataset_and_tokenizer(cfg)
|
| val_loader = DataLoader(
|
| val_ds,
|
| batch_size=cfg.train.batch_size,
|
| shuffle=False,
|
| collate_fn=lambda b: collate_conveyer_batch(b, pad_id=tokenizer.pad_id),
|
| )
|
|
|
| step_num = ckpt_meta.get("step", "?")
|
| elapsed_h = float(ckpt_meta.get("elapsed_sec", 0.0)) / 3600.0
|
|
|
| cryst = evaluate_crystallization_curve(model, cfg, val_loader, device, pad_id=tokenizer.pad_id)
|
| acc_c = cryst["per_slot_acc"]
|
| ent_c = cryst["per_slot_entropy"]
|
| B = len(acc_c)
|
|
|
| print("=" * 92)
|
| print(
|
| f"CONVEYOR-BELT DIFFUSION LLM BENCHMARK SUITE | Ckpt: {args.ckpt} (step {step_num}, {elapsed_h:.2f}h) "
|
| f"| Params: {model.num_parameters()/1e6:.2f}M | Device: {device}"
|
| )
|
| print("=" * 92)
|
| print(
|
| f"[1] Held-Out Real TinyStories LM Metrics (10,000 unseen stories, 0% templates):\n"
|
| f" - Conveyor Slot-0 Perplexity : {cryst['slot0_ppl']:6.2f} | Top-1 Next-Token Acc: {cryst['slot0_acc']*100:5.2f}%\n"
|
| f" - Causal Prefix LM Perplexity: {cryst['causal_lm_ppl']:6.2f} | Top-1 Next-Token Acc: {cryst['causal_lm_acc']*100:5.2f}%\n"
|
| f" - Back-75% Planning Pause : {cryst['mean_plan_pause']*100:5.1f}%"
|
| )
|
| print("-" * 92)
|
| print("[2] Per-Slot Crystallization Curve along the 64-Slot Oven Belt (Back k=63 -> Front k=0):")
|
| sample_slots = [B - 1, (3 * B) // 4, B // 2, B // 4, min(8, B - 1), min(3, B - 1), 0]
|
| seen: set[int] = set()
|
| for k in sample_slots:
|
| if k in seen:
|
| continue
|
| seen.add(k)
|
| zone = "Back (75% Plan) " if k >= 16 else "Front (25% Stream)"
|
| print(f" Slot {k:2d} [{zone}] | Top-1 Acc: {acc_c[k]*100:6.2f}% | Entropy: {ent_c[k]:.3f} nats")
|
|
|
| if cfg.data.dataset == "tinystories":
|
| cs = evaluate_zero_shot_commonsense(model, tokenizer, device)
|
| gen_q = evaluate_prefix_sensitivity_and_generation_quality(model, tokenizer, device)
|
| print("-" * 92)
|
| print(
|
| f"[3] Zero-Shot StoryCloze / Commonsense 2-Way Multiple Choice ({cs['num_pairs']} unseen pairs, random=50%):\n"
|
| f" - Causal LM Choice Accuracy : {cs['commonsense_lm_acc']*100:6.2f}% (mean log-prob margin: {cs['mean_logprob_margin']:+.3f} nats)\n"
|
| f" - Conveyor Slot-0 Choice Acc : {cs['commonsense_slot0_acc']*100:6.2f}%"
|
| )
|
| print("-" * 92)
|
| print(
|
| f"[4] Prompt Responsiveness & Open-Ended Generation Quality (6 novel story prompts):\n"
|
| f" - Word-Order Sensitivity KL : {gen_q['prefix_word_order_kl']:6.3f} nats (broken model was 0.032; >0.50 = strong)\n"
|
| f" - Prompt Entity Retention : {gen_q['prompt_entity_retention']*100:6.1f}%\n"
|
| f" - Lexical Diversity : Distinct-2 = {gen_q['distinct_2']*100:5.1f}% | Distinct-3 = {gen_q['distinct_3']*100:5.1f}%\n"
|
| f" - 4-gram Repetition Rate : {gen_q['repeat_4gram_rate']*100:5.2f}% (lower is better; 0% = zero loops)"
|
| )
|
| print("-" * 92)
|
| print("[5] Live Novel Prompt Continutations:")
|
| for s in gen_q["samples"][:3]:
|
| print(f" PROMPT : {s['prompt']}")
|
| print(f" OUTPUT : {s['continuation']}\n")
|
| else:
|
| cs = {}
|
| gen_q = {}
|
|
|
| print("-" * 92)
|
| print(f"{'Bake Depth K':<15} | {'Prompt Overlap W':<18} | {'Pos-0 Acc':<12} | {'Token Acc':<12} | {'Exact Seq Acc':<14}")
|
| print("-" * 92)
|
|
|
| sweep_results: list[dict[str, Any]] = []
|
| for k_bake in args.bake_steps:
|
| for w_ov in args.overlaps:
|
| res = evaluate_closed_loop_generation(
|
| model=model,
|
| val_ds=val_ds,
|
| tokenizer=tokenizer,
|
| device=device,
|
| num_episodes=args.episodes,
|
| prompt_overlap=w_ov,
|
| bake_steps=k_bake,
|
| )
|
| print(
|
| f"K = {k_bake:<11d} | W_pre = {w_ov:<10d} | "
|
| f"{res['first_token_acc']*100:9.2f}% | {res['token_acc']*100:9.2f}% | {res['exact_seq_acc']*100:10.2f}%"
|
| )
|
| sweep_results.append({"bake_steps": k_bake, "prompt_overlap": w_ov, **res})
|
|
|
| if args.out_json:
|
| Path(args.out_json).parent.mkdir(parents=True, exist_ok=True)
|
| Path(args.out_json).write_text(
|
| json.dumps(
|
| {
|
| "step": step_num,
|
| "crystallization": cryst,
|
| "commonsense": cs,
|
| "generation_quality": gen_q,
|
| "closed_loop": sweep_results,
|
| },
|
| indent=2,
|
| ),
|
| encoding="utf-8",
|
| )
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
