chat.py

#!/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()