#!/usr/bin/env python3 """ PC-SHO-DLM Chat — Conditional masked diffusion chatbot. The prompt stays unmasked. The response region starts fully masked. PC settling iteratively denoises the response conditioned on the prompt. Usage: python3 train_chat.py --data alpaca # Train on dialogue data python3 chat.py --checkpoint checkpoints/chat/chat_final.pt # Chat """ import argparse, glob, os, sys, torch import torch.nn.functional as F sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src")) from model import PCSHODLM, PCSHOConfig, LocalParameterUpdater def find_checkpoint(): for p in ["checkpoints/chat/chat_final.pt", "checkpoints/chat/*.pt", "checkpoints/*/checkpoint_final.pt", "checkpoints/*/*.pt"]: files = sorted(glob.glob(p), key=os.path.getmtime, reverse=True) if files: return files[0] return None def encode_text(text, vocab_size=257): return [min(b + 1, vocab_size - 1) for b in text.encode("utf-8")] def decode_tokens(tokens): raw = tokens if isinstance(tokens, list) else tokens.squeeze().cpu().tolist() # Filter out mask tokens (0) and padding return bytes([max(0, t - 1) for t in raw if t > 0 and t < 257]).decode("utf-8", errors="replace") def generate_response(model, prompt_text, response_len=256, device="cpu", temperature=0.8, n_steps=None): """Generate a response conditioned on prompt via conditional masked diffusion. 1. Encode prompt as bytes 2. Create: [prompt_tokens | MASK * response_len] 3. Reverse diffusion: iteratively unmask response tokens 4. PC settling at each step refines predictions conditioned on prompt """ config = model.config T = n_steps or config.n_diffusion_steps mask_id = config.mask_token_id prompt_tokens = encode_text(prompt_text) prompt_len = len(prompt_tokens) # Ensure we fit in max_seq_len max_response = config.max_seq_len - prompt_len if max_response <= 0: prompt_tokens = prompt_tokens[:config.max_seq_len // 2] prompt_len = len(prompt_tokens) max_response = config.max_seq_len - prompt_len response_len = min(response_len, max_response) # Build initial sequence: [prompt | MASK...] seq = prompt_tokens + [mask_id] * response_len if len(seq) < config.max_seq_len: seq = seq + [mask_id] * (config.max_seq_len - len(seq)) x = torch.tensor(seq[:config.max_seq_len], dtype=torch.long).unsqueeze(0).to(device) # Response region mask (only unmask response, never touch prompt) response_region = torch.zeros(1, config.max_seq_len, dtype=torch.bool, device=device) response_region[0, prompt_len:prompt_len + response_len] = True prev_h = None prev_v = None with torch.no_grad(): for t_val in range(T, 0, -1): t = torch.full((1,), t_val, dtype=torch.long, device=device) # Only consider masked tokens in response region mask = (x == mask_id) & response_region if not mask.any(): break # Embed and settle h_0 = model.embed_input(x, t) h_init = model.amortized_forward_pass(h_0) h_settled, v_final, _, _, _ = model.settle( h_init, x, mask, t, prev_h=prev_h, prev_v=prev_v, ) # State continuation prev_h = h_settled prev_v = v_final # Decode with temperature logits = model.readout(model.readout_norm(h_settled[-1])) if temperature > 0: probs = F.softmax(logits / temperature, dim=-1) else: probs = F.softmax(logits, dim=-1) # Determine how many tokens to unmask t_prev = t_val - 1 target_mask_rate = model.schedule.mask_rate( torch.tensor([t_prev], device=device) ).squeeze() n_currently_masked = mask.sum().item() n_to_keep_masked = int(target_mask_rate * response_len) n_to_unmask = max(1, n_currently_masked - n_to_keep_masked) # Score by confidence, only in response region confidence = probs.max(dim=-1).values.squeeze() # (S,) confidence[~mask.squeeze()] = float("inf") # Unmask the most confident response tokens masked_indices = mask.squeeze().nonzero(as_tuple=True)[0] if len(masked_indices) > 0: conf_at_masked = confidence[masked_indices] n_to_unmask = min(n_to_unmask, len(masked_indices)) _, top_idx = conf_at_masked.topk(n_to_unmask) unmask_pos = masked_indices[top_idx] # Sample tokens if temperature > 0: for pos in unmask_pos: token = torch.multinomial(probs[0, pos], 1).item() x[0, pos] = token else: x[0, unmask_pos] = logits[0, unmask_pos].argmax(dim=-1) # Final pass: unmask any remaining remaining = (x == mask_id) & response_region if remaining.any(): t_final = torch.ones(1, dtype=torch.long, device=device) h_0 = model.embed_input(x, t_final) h_init = model.amortized_forward_pass(h_0) h_s, _, _, _, _ = model.settle(h_init, x, remaining, t_final, prev_h=prev_h, prev_v=prev_v) logits = model.readout(model.readout_norm(h_s[-1])) if temperature > 0: for pos in remaining.squeeze().nonzero(as_tuple=True)[0]: probs = F.softmax(logits[0, pos] / temperature, dim=-1) x[0, pos] = torch.multinomial(probs, 1).item() else: x[0, remaining.squeeze()] = logits[0, remaining.squeeze()].argmax(dim=-1) # Extract response response_tokens = x[0, prompt_len:prompt_len + response_len].cpu().tolist() return decode_tokens(response_tokens) def main(): parser = argparse.ArgumentParser(description="PC-SHO-DLM Chat") parser.add_argument("--checkpoint", type=str, default=None) parser.add_argument("--device", type=str, default="auto") parser.add_argument("--temperature", type=float, default=0.8) parser.add_argument("--response_len", type=int, default=256) args = parser.parse_args() if args.device == "auto": if torch.cuda.is_available(): args.device = "cuda" elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): args.device = "mps" else: args.device = "cpu" ckpt_path = args.checkpoint or find_checkpoint() if not ckpt_path or not os.path.exists(ckpt_path): print("No checkpoint found. Train first:") print(" python3 train_chat.py --data alpaca") sys.exit(1) print(f"Loading {ckpt_path}...") ckpt = torch.load(ckpt_path, map_location=args.device, weights_only=False) config = ckpt["config"] model = PCSHODLM(config).to(args.device) model.load_state_dict(ckpt["model_state_dict"], strict=False) model.eval() print(f"Model: {sum(p.numel() for p in model.parameters()):,} params on {args.device}") print() print("=" * 60) print(" PC-SHO-DLM Chat") print(" Conditional Masked Diffusion Chatbot") print("=" * 60) print() print(" Type a message and the model generates a response.") print(" The prompt stays fixed; the response is iteratively") print(" denoised via predictive-coding energy minimization.") print() print(" /temp 0.5 Set temperature (0=greedy, 1=creative)") print(" /len 512 Set response length") print(" /quit Exit") print() temp = args.temperature resp_len = args.response_len while True: try: user_input = input("You> ").strip() except (EOFError, KeyboardInterrupt): print("\nBye!") break if not user_input: continue if user_input.lower() in ["/quit", "/exit", "/q"]: break elif user_input.lower().startswith("/temp"): try: temp = float(user_input.split()[1]) except: pass print(f" Temperature: {temp}") continue elif user_input.lower().startswith("/len"): try: resp_len = int(user_input.split()[1]) except: pass print(f" Response length: {resp_len}") continue print(" [Generating...]") response = generate_response(model, user_input, response_len=resp_len, device=args.device, temperature=temp) # Clean up: strip trailing nulls/garbage response = response.rstrip('\x00').strip() print(f"\nBot> {response}\n") if __name__ == "__main__": main()