train_chat.py

"""
PC-SHO-DLM Chat Training — Conditional Masked Diffusion

Trains the model to generate responses conditioned on prompts.
Prompt stays unmasked, response gets masked/denoised via PC settling.
"""

import argparse, json, os, sys, time
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
from model import PCSHODLM, PCSHOConfig, LocalParameterUpdater, count_parameters


class DialogueDataset(Dataset):
    """Byte-level dialogue dataset. Each sample = (prompt, response) pair.

    Returns input_ids (full sequence) + prompt_len (where response starts).
    During training, only the response portion gets masked.
    """

    def __init__(self, pairs: list, max_seq_len: int, vocab_size: int = 257):
        self.max_seq_len = max_seq_len
        self.vocab_size = vocab_size
        self.samples = []

        sep = [0x0A, 0x0A]  # double newline as separator

        for prompt, response in pairs:
            p_bytes = [min(b + 1, vocab_size - 1) for b in prompt.encode("utf-8")]
            r_bytes = [min(b + 1, vocab_size - 1) for b in response.encode("utf-8")]
            s_bytes = [min(b + 1, vocab_size - 1) for b in bytes(sep)]

            # Truncate: prompt gets up to 40% of seq_len, response gets the rest
            max_p = int(max_seq_len * 0.4)
            max_r = max_seq_len - min(len(p_bytes), max_p) - len(s_bytes)

            p = p_bytes[:max_p]
            r = r_bytes[:max_r]

            full = p + s_bytes + r
            prompt_len = len(p) + len(s_bytes)

            # Pad
            if len(full) < max_seq_len:
                full = full + [0] * (max_seq_len - len(full))  # 0 = MASK token as pad
            full = full[:max_seq_len]

            self.samples.append({
                "input_ids": torch.tensor(full, dtype=torch.long),
                "prompt_len": prompt_len,
                "response_len": len(r),
            })

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx]


def load_dialogue_data(name="alpaca", max_samples=None):
    """Load dialogue pairs from HuggingFace."""
    from datasets import load_dataset

    pairs = []

    if name == "alpaca":
        print("Loading Alpaca dataset...")
        try:
            ds = load_dataset("tatsu-lab/alpaca", split="train")
            for row in ds:
                instruction = row.get("instruction", "")
                inp = row.get("input", "")
                output = row.get("output", "")
                prompt = f"{instruction}\n{inp}".strip() if inp else instruction
                if prompt and output:
                    pairs.append((prompt, output))
                    if max_samples and len(pairs) >= max_samples:
                        break
        except Exception as e:
            print(f"HF load failed ({e}), downloading JSON directly...")
            import urllib.request, json as json_mod
            url = "https://raw.githubusercontent.com/tatsu-lab/stanford_alpaca/main/alpaca_data.json"
            with urllib.request.urlopen(url) as r:
                data = json_mod.loads(r.read())
            for row in data:
                instruction = row.get("instruction", "")
                inp = row.get("input", "")
                output = row.get("output", "")
                prompt = f"{instruction}\n{inp}".strip() if inp else instruction
                if prompt and output:
                    pairs.append((prompt, output))
                    if max_samples and len(pairs) >= max_samples:
                        break

    elif name == "dolly":
        print("Loading Dolly dataset...")
        ds = load_dataset("databricks/dolly-15k", split="train")
        for row in ds:
            prompt = row.get("instruction", "") or row.get("context", "")
            response = row.get("response", "")
            if prompt and response:
                pairs.append((prompt, response))
                if max_samples and len(pairs) >= max_samples:
                    break

    elif name == "oasst":
        print("Loading OpenAssistant dataset...")
        ds = load_dataset("OpenAssistant/oasst1", split="train")
        # Flatten conversation trees into (parent, reply) pairs
        messages = {row["message_id"]: row for row in ds}
        for row in ds:
            if row["parent_id"] and row["parent_id"] in messages:
                parent = messages[row["parent_id"]]
                pairs.append((parent["text"], row["text"]))
                if max_samples and len(pairs) >= max_samples:
                    break

    else:
        raise ValueError(f"Unknown dataset: {name}")

    print(f"Loaded {len(pairs)} dialogue pairs")
    return pairs


def conditional_corrupt(x_0, t, mask_token_id, schedule, prompt_lens):
    """Mask only response tokens, keep prompt unmasked."""
    B, S = x_0.shape
    mask_prob = schedule.mask_rate(t).unsqueeze(-1)  # (B, 1)
    random_mask = torch.rand_like(x_0.float()) < mask_prob  # (B, S)

    # Only mask response positions
    response_mask = torch.zeros_like(x_0, dtype=torch.bool)
    for b in range(B):
        response_mask[b, prompt_lens[b]:] = True

    mask = random_mask & response_mask
    x_t = x_0.clone()
    x_t[mask] = mask_token_id
    return x_t, mask


def train():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", default="alpaca", choices=["alpaca", "dolly", "oasst"])
    parser.add_argument("--max_samples", type=int, default=None)
    parser.add_argument("--d_model", type=int, default=256)
    parser.add_argument("--n_layers", type=int, default=6)
    parser.add_argument("--n_heads", type=int, default=8)
    parser.add_argument("--seq_len", type=int, default=512)
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--max_steps", type=int, default=5000)
    parser.add_argument("--device", type=str, default="auto")
    parser.add_argument("--save_dir", type=str, default="checkpoints/chat")
    parser.add_argument("--checkpoint", type=str, default=None, help="Fine-tune from existing checkpoint")
    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"

    # Config
    config = PCSHOConfig(
        vocab_size=257, max_seq_len=args.seq_len,
        d_model=args.d_model, n_heads=args.n_heads, n_layers=args.n_layers,
        d_ff=args.d_model * 4, n_diffusion_steps=64, n_settling_steps=6,
        mask_token_id=0, dropout=0.1, feedback_rank=args.d_model // 4,
        denoising_lambda=500.0,
        homeostatic_target=0.0, lateral_inhibition=0.0,
    )

    # Load data
    pairs = load_dialogue_data(args.data, args.max_samples)
    n_val = max(1, len(pairs) // 10)
    train_pairs = pairs[n_val:]
    val_pairs = pairs[:n_val]

    train_ds = DialogueDataset(train_pairs, config.max_seq_len)
    val_ds = DialogueDataset(val_pairs, config.max_seq_len)

    print(f"Train: {len(train_ds)}, Val: {len(val_ds)}")

    # Model
    model = PCSHODLM(config).to(args.device)

    if args.checkpoint:
        print(f"Loading checkpoint: {args.checkpoint}")
        ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=False)
        model.load_state_dict(ckpt["model_state_dict"], strict=False)

    model.train()
    n_params = count_parameters(model)
    print(f"Model: {n_params:,} params on {args.device}")

    # Training
    updater = LocalParameterUpdater(model, lr_forward=args.lr, lr_feedback=args.lr,
                                     lr_readout=args.lr, lr_precision=args.lr * 0.1)

    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True)
    val_loader = DataLoader(val_ds, batch_size=args.batch_size)

    os.makedirs(args.save_dir, exist_ok=True)
    step = 0
    start = time.time()
    K_max = config.n_settling_steps

    print(f"\n{'='*60}")
    print(f"Chat Training | {n_params:,} params | {args.data} | {len(train_ds)} samples")
    print(f"{'='*60}\n")

    while step < args.max_steps:
        for batch in train_loader:
            if step >= args.max_steps:
                break

            x_0 = batch["input_ids"].to(args.device)
            prompt_lens = batch["prompt_len"]
            B, S = x_0.shape

            # Curriculum K
            K_curr = max(3, int(K_max * min(1.0, step / 1000)))
            model.config.n_settling_steps = K_curr

            # Sample timestep
            t = torch.randint(1, config.n_diffusion_steps + 1, (B,), device=args.device)

            # Conditional corruption — only mask response
            x_t, mask = conditional_corrupt(x_0, t, config.mask_token_id, model.schedule, prompt_lens)

            # HONEST SETTLE: use x_t (corrupted), NOT x_0 (answer)
            # Settling minimizes entropy (push toward confident predictions)
            # without peeking at the ground truth
            h_0 = model.embed_input(x_t, t)
            h_init = model.amortized_forward_pass(h_0)

            with torch.no_grad():
                h_settled, _, energies, _, _ = model.settle(h_init, x_t, mask, t)

            # Local updates (same as LocalParameterUpdater but with our custom mask)
            L = config.n_layers
            for l in range(L):
                h_below = h_settled[l].detach()
                h_current = h_settled[l + 1].detach()
                updater.forward_optimizers[l].zero_grad()
                mu_up = model.forward_blocks[l](h_below)
                with torch.no_grad():
                    tok_p, ch_p = model.precision_up[l](h_current, t)
                eps = h_current - mu_up
                loss_l = 0.5 * (eps * tok_p * ch_p * eps).sum() / B
                loss_l.backward()
                updater.forward_optimizers[l].step()

            for l in range(L - 1):
                updater.feedback_optimizers[l].zero_grad()
                h_above = h_settled[l + 2].detach()
                h_current = h_settled[l + 1].detach()
                mu_down = model.feedback_blocks[l](h_above)
                with torch.no_grad():
                    tok_p, ch_p = model.precision_down[l](h_current, t)
                eps = h_current - mu_down
                loss_l = 0.5 * (eps * tok_p * ch_p * eps).sum() / B
                loss_l.backward()
                updater.feedback_optimizers[l].step()

            # Readout
            updater.readout_optimizer.zero_grad()
            logits = model.readout(model.readout_norm(h_settled[L].detach()))
            mask_logits = logits[mask]
            mask_targets = x_0[mask]
            if mask_logits.numel() > 0:
                readout_loss = F.cross_entropy(mask_logits, mask_targets)
                readout_loss.backward()
                updater.readout_optimizer.step()
                loss_val = readout_loss.item()
            else:
                loss_val = 0.0

            # Embeddings
            updater.embed_optimizer.zero_grad()
            with torch.enable_grad():
                h_0_fresh = model.embed_input(x_t, t)
                h_init_fresh = model.amortized_forward_pass(h_0_fresh)
                logits_e = model.readout(model.readout_norm(h_init_fresh[L]))
                ml_e, mt_e = logits_e[mask], x_0[mask]
                if ml_e.numel() > 0:
                    embed_loss = F.cross_entropy(ml_e, mt_e)
                    embed_loss.backward()
                    torch.nn.utils.clip_grad_norm_(
                        list(model.token_embed.parameters()) + list(model.pos_embed.parameters()) + list(model.time_embed.parameters()), 1.0
                    )
                    updater.embed_optimizer.step()

            step += 1

            if step % 50 == 0:
                el = time.time() - start
                e = f"{energies[-1]:.0f}" if energies else "N/A"
                print(f"Step {step:5d} | Loss: {loss_val:.4f} | Energy: {e} | K={K_curr} | {el:.0f}s")

            if step % 500 == 0:
                torch.save({
                    "model_state_dict": model.state_dict(),
                    "config": config,
                    "step": step,
                }, os.path.join(args.save_dir, f"chat_step{step}.pt"))

    # Final save
    torch.save({
        "model_state_dict": model.state_dict(),
        "config": config,
        "step": step,
    }, os.path.join(args.save_dir, "chat_final.pt"))

    print(f"\nDone! {step} steps in {time.time()-start:.0f}s")
    print(f"Checkpoint: {args.save_dir}/chat_final.pt")
    print(f"Run: python3 chat.py --checkpoint {args.save_dir}/chat_final.pt")


if __name__ == "__main__":
    train()