| from __future__ import annotations
|
|
|
| import argparse
|
| import sys
|
| import time
|
| import torch
|
|
|
| from conveyer.checkpoint import load_checkpoint
|
| from conveyer.config import load_config
|
| from conveyer.data import build_dataset_and_tokenizer
|
| from conveyer.infer import TickSnapshot, stream_conveyer_generation
|
| from conveyer.model import ConveyerDiffusionLM
|
|
|
| RESET = "\033[0m"
|
| BOLD = "\033[1m"
|
| DIM = "\033[2m"
|
| GREEN = "\033[92m"
|
| CYAN = "\033[96m"
|
| YELLOW = "\033[93m"
|
| MAGENTA = "\033[95m"
|
| GRAY = "\033[90m"
|
|
|
|
|
| def _format_slot_badge(slot_idx: int, piece: str, prob: float) -> str:
|
| clean = piece.replace("\n", " ")[:9]
|
| if slot_idx == 0:
|
| return f"{BOLD}{GREEN}[0:{clean} {prob*100:2.0f}%]{RESET}"
|
| elif slot_idx < 15:
|
| return f"{GREEN}{clean}({prob*100:2.0f}%){RESET}"
|
| elif slot_idx < 40:
|
| return f"{YELLOW}{clean}({prob*100:2.0f}%){RESET}"
|
| else:
|
| return f"{GRAY}{clean}{RESET}"
|
|
|
|
|
| def render_frame(
|
| prompt_str: str,
|
| committed_pieces: list[str],
|
| snap: TickSnapshot,
|
| tokenizer: Any,
|
| clear_screen: bool = True,
|
| ) -> str:
|
| B = len(snap.slot_top1_tokens)
|
| lines: list[str] = []
|
| if clear_screen:
|
| lines.append("\033[2J\033[H")
|
| lines.append("=" * 98)
|
| mode_tag = (
|
| f"{CYAN}PRE-BAKING WARMUP (Clamping Input -> Preparing Conveyor Belt){RESET}"
|
| if snap.is_warmup
|
| else f"{GREEN}LIVE CONVEYOR BAKING (Serving Slot 0 -> Shifting 1..63 -> Spawning 63){RESET}"
|
| )
|
| lines.append(f"{BOLD}64-SLOT DUAL-CHANNEL CONVEYOR OVEN LLM{RESET} | Tick t={snap.tick:3d} | {mode_tag}")
|
| lines.append("-" * 98)
|
| lines.append(f"{BOLD}Input Prompt :{RESET} {prompt_str}")
|
| committed_str = " ".join(committed_pieces) if committed_pieces else "(warming up belt...)"
|
| lines.append(f"{BOLD}Served Output:{RESET} {BOLD}{GREEN}{committed_str}{RESET}")
|
| lines.append("-" * 98)
|
|
|
|
|
| front_items = [
|
| _format_slot_badge(k, tokenizer.id_to_piece(snap.slot_top1_tokens[k]), snap.slot_top1_probs[k])
|
| for k in range(min(15, B))
|
| ]
|
| lines.append(f"{BOLD}{GREEN}FRONT [00..14] (Final Polish, alpha->0):{RESET}")
|
| lines.append(" <- POP " + " | ".join(front_items[:8]))
|
| if len(front_items) > 8:
|
| lines.append(" " + " | ".join(front_items[8:15]))
|
|
|
|
|
| if B > 15:
|
| mid_indices = list(range(15, min(40, B), 2))
|
| mid_items = [
|
| _format_slot_badge(k, tokenizer.id_to_piece(snap.slot_top1_tokens[k]), snap.slot_top1_probs[k])
|
| for k in mid_indices[:10]
|
| ]
|
| lines.append(f"{BOLD}{YELLOW}MIDDLE [15..39] (Firming Phrasing, Dual-Channel):{RESET}")
|
| lines.append(" " + " | ".join(mid_items))
|
|
|
|
|
| if B > 40:
|
| back_indices = list(range(40, B, 2))
|
| if (B - 1) not in back_indices:
|
| back_indices.append(B - 1)
|
| back_items = [
|
| f"[{k}:{_format_slot_badge(k, tokenizer.id_to_piece(snap.slot_top1_tokens[k]), snap.slot_top1_probs[k])}]"
|
| for k in back_indices[-8:]
|
| ]
|
| lines.append(f"{BOLD}{MAGENTA}BACK [40..63] (Rough Ideas <- Slot 63 Raw Dough Spawn):{RESET}")
|
| lines.append(" " + " ".join(back_items) + " <- [NEW DOUGH]")
|
|
|
|
|
| ent_front = sum(snap.slot_entropies[: min(8, B)]) / max(1, min(8, B))
|
| ent_mid = sum(snap.slot_entropies[min(15, B - 1) : min(40, B)]) / max(1, min(40, B) - min(15, B - 1))
|
| ent_back = sum(snap.slot_entropies[max(0, B - 16) : B]) / max(1, B - max(0, B - 16))
|
| gate_avg = sum(snap.slot_gate_means) / max(1, B)
|
| lines.append("-" * 98)
|
| lines.append(
|
| f"Entropy (Front/Mid/Back): {ent_front:.2f} / {ent_mid:.2f} / {ent_back:.2f} nats | "
|
| f"Mean Thought Gate g_k: {gate_avg:.3f}"
|
| )
|
| lines.append("=" * 98)
|
| return "\n".join(lines)
|
|
|
|
|
| def main() -> None:
|
| parser = argparse.ArgumentParser(description="Live Terminal Visualizer for the 64-Slot Conveyor-Belt Diffusion LLM")
|
| parser.add_argument("--ckpt", type=str, default=None, help="Checkpoint path (.pt). If omitted, initializes from config")
|
| parser.add_argument("--config", type=str, default="configs/synthetic.yaml", help="Config path if --ckpt is not provided")
|
| parser.add_argument("--prompt", type=str, default=None, help="Custom input prompt")
|
| parser.add_argument("--tokens", type=int, default=24, help="Number of output tokens to bake")
|
| parser.add_argument("--overlap", type=int, default=12, help="Prompt-overlap pre-baking window (n - W)")
|
| parser.add_argument("--bake-steps", type=int, default=64, help="Conveyor belt size (1..64)")
|
| parser.add_argument("--hz", type=float, default=10.0, help="Refresh rate in ticks/sec (use 0 for no delay)")
|
| parser.add_argument("--no-clear", action="store_true", help="Print frames sequentially without ANSI screen clear")
|
| args = parser.parse_args()
|
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| if args.ckpt is not None:
|
| model, cfg, _ = load_checkpoint(args.ckpt, device=device)
|
| else:
|
| cfg = load_config(args.config)
|
| model = ConveyerDiffusionLM(cfg.model).to(device)
|
|
|
| _, val_ds, tokenizer = build_dataset_and_tokenizer(cfg)
|
|
|
| if args.prompt is not None:
|
| prompt_ids = tokenizer.encode(args.prompt)
|
| prompt_tensor = torch.tensor([prompt_ids], dtype=torch.long, device=device)
|
| else:
|
| sample = val_ds[0]
|
| p_len = int(sample["prompt_len"])
|
| prompt_tensor = sample["input_ids"][:p_len].unsqueeze(0).to(device)
|
|
|
| prompt_str = tokenizer.decode(prompt_tensor[0].tolist())
|
| committed: list[str] = []
|
| delay = (1.0 / args.hz) if args.hz > 0 else 0.0
|
|
|
| for tok, snap in stream_conveyer_generation(
|
| model=model,
|
| prompt_tokens=prompt_tensor,
|
| max_new_tokens=args.tokens,
|
| prompt_overlap=args.overlap,
|
| bake_steps=args.bake_steps,
|
| temperature=0.0,
|
| eos_token_id=tokenizer.eos_id,
|
| ):
|
| if not snap.is_warmup:
|
| committed.append(tokenizer.id_to_piece(int(tok[0].item())))
|
| frame = render_frame(
|
| prompt_str=prompt_str,
|
| committed_pieces=committed,
|
| snap=snap,
|
| tokenizer=tokenizer,
|
| clear_screen=(not args.no_clear),
|
| )
|
| sys.stdout.write(frame + "\n")
|
| sys.stdout.flush()
|
| if delay > 0:
|
| time.sleep(delay)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
