"""
class ConveyerWebApp:
def __init__(
self,
model: ConveyerDiffusionLM,
cfg: Any,
ckpt_path: str,
ckpt_meta: dict[str, Any],
val_ds: Any,
tokenizer: Any,
device: torch.device,
) -> None:
self.model = model
self.cfg = cfg
self.ckpt_path = ckpt_path
self.ckpt_meta = ckpt_meta
self.val_ds = val_ds
self.tokenizer = tokenizer
self.device = device
self.cached_cryst = self._compute_initial_stats()
def _compute_initial_stats(self) -> list[dict[str, Any]]:
metrics = self.ckpt_meta.get("metrics", {})
acc_curve = metrics.get("val_per_slot_acc")
ent_curve = metrics.get("val_per_slot_entropy")
if not acc_curve or not ent_curve:
val_loader = DataLoader(
self.val_ds,
batch_size=min(32, len(self.val_ds)),
shuffle=False,
collate_fn=lambda b: collate_conveyer_batch(b, pad_id=self.tokenizer.pad_id),
)
c = evaluate_crystallization_curve(
self.model, self.cfg, val_loader, self.device, pad_id=self.tokenizer.pad_id, max_batches=2
)
acc_curve = c["per_slot_acc"]
ent_curve = c["per_slot_entropy"]
B = len(acc_curve)
slots_to_show = [0, 3, 8, 15, 24, 36, 48, B - 1]
rows: list[dict[str, Any]] = []
seen: set[int] = set()
for k in slots_to_show:
k = min(k, B - 1)
if k in seen:
continue
seen.add(k)
zone = (
"Slot 0 Pop"
if k == 0
else ("Front Conveyor" if k < 16 else ("Mid Expansion" if k == 15 else "Back 75% Plan"))
)
rows.append(
{
"slot": k,
"zone": zone,
"acc": float(acc_curve[k]),
"entropy": float(ent_curve[k]),
}
)
return rows
def compute_trace(
self,
prompt_text: str | None,
sample_idx: int,
overlap: int,
bake_steps: int,
max_tokens: int,
temperature: float = 0.0,
) -> dict[str, Any]:
gt_continuation: str | None = None
gt_target_ids: list[int] = []
if prompt_text is not None and prompt_text.strip() != "":
encoded = self.tokenizer.encode(prompt_text.strip())
prompt_ids = torch.tensor([encoded], dtype=torch.long, device=self.device)
for i in range(min(300, len(self.val_ds))):
seq = self.val_ds[i]["input_ids"]
if seq[: len(encoded)].tolist() == encoded:
rem = [
int(t)
for t in seq[len(encoded) :].tolist()
if int(t) not in (self.tokenizer.pad_id, self.tokenizer.eos_id)
]
gt_target_ids = rem[:max_tokens]
gt_continuation = self.tokenizer.decode(gt_target_ids)
break
else:
idx = max(0, min(sample_idx, len(self.val_ds) - 1))
item = self.val_ds[idx]
plen = int(item["prompt_len"])
seq = item["input_ids"]
prompt_ids = seq[:plen].unsqueeze(0).to(self.device)
rem = [int(t) for t in seq[plen:].tolist() if int(t) not in (self.tokenizer.pad_id, self.tokenizer.eos_id)]
gt_target_ids = rem[:max_tokens]
gt_continuation = self.tokenizer.decode(gt_target_ids)
full_out, snapshots = generate_conveyer(
model=self.model,
prompt_tokens=prompt_ids,
max_new_tokens=max_tokens,
prompt_overlap=overlap,
bake_steps=bake_steps,
temperature=temperature,
eos_token_id=self.tokenizer.eos_id,
deterministic_spawn=(temperature <= 1e-5),
)
plen_actual = prompt_ids.size(1)
gen_ids = full_out[0, plen_actual:].tolist()
sample_token_acc: float | None = None
if gt_target_ids:
comp_len = min(len(gen_ids), len(gt_target_ids))
if comp_len > 0:
matches = sum(1 for a, b in zip(gen_ids[:comp_len], gt_target_ids[:comp_len]) if a == b)
sample_token_acc = matches / comp_len
F_slots = min(self.model.front_slots, bake_steps)
serialized_snaps: list[dict[str, Any]] = []
for s in snapshots:
slots_info = []
B = len(s.slot_top1_tokens)
for k in range(B):
top3_list = [
{"w": self.tokenizer.id_to_piece(tid), "p": round(float(pr), 4)}
for tid, pr in zip(s.slot_top3_tokens[k], s.slot_top3_probs[k])
]
slots_info.append(
{
"k": k,
"piece": self.tokenizer.id_to_piece(s.slot_top1_tokens[k]),
"prob": round(float(s.slot_top1_probs[k]), 4),
"entropy": round(float(s.slot_entropies[k]), 3),
"gate": round(float(s.slot_gate_means[k]), 3),
"shift_prob": round(float(s.slot_shift_probs[k]), 3),
"top3": top3_list,
}
)
back_pause = (
100.0 * (1.0 - sum(s.slot_shift_probs[F_slots:]) / max(1, B - F_slots))
if B > F_slots
else 0.0
)
serialized_snaps.append(
{
"tick": s.tick,
"is_warmup": s.is_warmup,
"emitted_piece": self.tokenizer.id_to_piece(s.emitted_token),
"mean_gate": sum(s.slot_gate_means) / max(1, B),
"back_pause_pct": back_pause,
"slots": slots_info,
}
)
metrics = self.ckpt_meta.get("metrics", {})
val_s0_acc = float(metrics.get("val_slot0_acc", self.cached_cryst[0]["acc"] if self.cached_cryst else 0.0))
return {
"ckpt_name": self.ckpt_path,
"dataset": self.cfg.data.dataset,
"vocab_size": self.tokenizer.vocab_size,
"params_m": f"{self.model.num_parameters() / 1e6:.2f}",
"val_slot0_acc": val_s0_acc,
"sample_token_acc": sample_token_acc,
"prompt": self.tokenizer.decode(prompt_ids[0].tolist()),
"ground_truth_continuation": gt_continuation,
"crystallization": self.cached_cryst,
"snapshots": serialized_snaps,
}
def main() -> None:
parser = argparse.ArgumentParser(description="Interactive Web Simulator for the 64-Slot Conveyor-Belt Diffusion LLM")
parser.add_argument("--ckpt", type=str, default=None, help="Path to checkpoint (.pt)")
parser.add_argument("--config", type=str, default="configs/tinystories.yaml", help="Fallback config if no checkpoint")
parser.add_argument("--host", type=str, default=os.environ.get("HOST", "127.0.0.1"))
parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", "8766")))
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
ckpt_path = args.ckpt
if ckpt_path is None:
if Path("runs/tinystories/best.pt").exists():
ckpt_path = "runs/tinystories/best.pt"
elif Path("runs/synthetic/best.pt").exists():
ckpt_path = "runs/synthetic/best.pt"
else:
ckpt_path = "runs/tinystories/best.pt"
try:
model, cfg, ckpt_meta = load_checkpoint(ckpt_path, device=device)
except Exception:
cfg = load_config(args.config)
model = ConveyerDiffusionLM(cfg.model).to(device)
ckpt_path = "untrained"
ckpt_meta = {}
_, val_ds, tokenizer = build_dataset_and_tokenizer(cfg, infer_only=True)
app = ConveyerWebApp(
model=model,
cfg=cfg,
ckpt_path=ckpt_path,
ckpt_meta=ckpt_meta,
val_ds=val_ds,
tokenizer=tokenizer,
device=device,
)
class Handler(BaseHTTPRequestHandler):
def do_GET(self) -> None:
parsed = urlparse(self.path)
if parsed.path == "/api/trace":
qs = parse_qs(parsed.query)
prompt_text = qs.get("prompt", [None])[0]
sample_idx = int(qs.get("sample", ["0"])[0])
overlap = int(qs.get("overlap", ["12"])[0])
bake_steps = int(qs.get("bake_steps", ["64"])[0])
tokens = int(qs.get("tokens", ["32"])[0])
temp = float(qs.get("temp", ["0.0"])[0])
payload = json.dumps(
app.compute_trace(prompt_text, sample_idx, overlap, bake_steps, tokens, temperature=temp)
).encode("utf-8")
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Cache-Control", "no-store, no-cache, must-revalidate")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
else:
body = HTML_PAGE.encode("utf-8")
self.send_response(200)
self.send_header("Content-Type", "text/html; charset=utf-8")
self.send_header("Cache-Control", "no-store")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, format: str, *args: Any) -> None:
return
server = HTTPServer((args.host, args.port), Handler)
print(f"Conveyor-Belt Diffusion LLM Web Simulator running at http://{args.host}:{args.port}", flush=True)
server.serve_forever()
if __name__ == "__main__":
main()