ritikraj2425/Discrete-Diffusion-Text-Demo

1import torch2import torch.nn.functional as F3from tokenizers import Tokenizer4from train import MaskedDiffusionModel5from tokenizer import MAX_SEQ_LENGTH6import os7 8 9def load_tokenizer_full(vocab_file="subword_tokenizer.json"):10    tokenizer = Tokenizer.from_file(vocab_file)11    vocab     = tokenizer.get_vocab()12    id2word   = {int(v): k for k, v in vocab.items()}13    return tokenizer, vocab, id2word14 15 16def decode_response(token_ids, tokenizer):17    """18    Use the tokenizer's built-in decode() for proper BPE subword handling.19    Filter out special tokens before decoding.20    """21    special_ids = {22        tokenizer.token_to_id("[PAD]"),23        tokenizer.token_to_id("[BOS]"),24        tokenizer.token_to_id("[EOS]"),25        tokenizer.token_to_id("[MASK]"),26        tokenizer.token_to_id("[UNK]"),27    }28 29    filtered_ids = [tid for tid in token_ids if tid not in special_ids]30 31    if not filtered_ids:32        return ""33 34    # Let the tokenizer handle subword reassembly35    text = tokenizer.decode(filtered_ids)36 37    # Light cleanup38    text = text.strip()39    # Fix spacing around punctuation40    for p in [".", ",", "?", "!", "'", ":"]:41        text = text.replace(f" {p}", p)42 43    return text44 45 46def generate_response(47    model, tokenizer, id2word,48    prompt,49    max_response_length=24,50    sampling_steps=40,51    temperature=0.5,52    top_k=1553):54    """55    Timestep-aware iterative denoising with running confidence remasking.56 57    Key differences from old version:58    1. Passes timestep t to model at each step (model now knows what stage it's at)59    2. t = fraction of tokens still masked (starts ~1.0, ends ~0.0)60    3. Running confidence remasking: tracks cumulative confidence per token61    """62    model.eval()63    device = next(model.parameters()).device64 65    bos_id  = tokenizer.token_to_id("[BOS]")66    eos_id  = tokenizer.token_to_id("[EOS]")67    mask_id = tokenizer.token_to_id("[MASK]")68    pad_id  = tokenizer.token_to_id("[PAD]")69 70    formatted = f"user: {prompt.lower().strip()} bot:"71    input_ids = tokenizer.encode(formatted).ids72 73    # Clamp response length to available space74    max_resp = min(max_response_length, MAX_SEQ_LENGTH - len(input_ids) - 2)75    if max_resp <= 0:76        print("Prompt too long.")77        return ""78 79    sequence = [bos_id] + input_ids + [mask_id] * max_resp + [eos_id]80    sequence += [pad_id] * (MAX_SEQ_LENGTH - len(sequence))81    seq_tensor = torch.tensor([sequence], dtype=torch.long, device=device)82 83    response_start = 1 + len(input_ids)84    response_end   = response_start + max_resp85    mask_indices   = list(range(response_start, response_end))86    num_masks      = len(mask_indices)87 88    # Running confidence: tracks cumulative confidence per token position89    # Tokens that are consistently predicted with high confidence get revealed first90    running_confidence = torch.zeros(num_masks, device=device)91 92    for step in range(1, sampling_steps + 1):93        # ── Timestep: fraction of tokens still masked ──94        # Step 1: t ≈ 1.0 (most tokens masked, early stage)95        # Step 40: t ≈ 0.025 (few tokens masked, final refinement)96        # Clamp to [0.05, 1.0] to stay within training range97        t_val = max(1.0 - step / sampling_steps, 0.05)98        t = torch.tensor([t_val], device=device)99 100        # Padding mask so attention ignores PAD positions101        pad_mask = (seq_tensor == pad_id)102 103        with torch.no_grad():104            logits = model(seq_tensor, t, src_key_padding_mask=pad_mask)105 106        # Top-k filtering — zero out all but top-k logits107        response_logits = logits[0, mask_indices]  # [num_masks, vocab_size]108        if top_k > 0:109            top_k_vals, _ = torch.topk(response_logits, top_k, dim=-1)110            min_top_k     = top_k_vals[:, -1].unsqueeze(-1)111            response_logits = response_logits.masked_fill(response_logits < min_top_k, float('-inf'))112 113        # Temperature scaling and sampling114        scaled = response_logits / max(temperature, 1e-6)115        probs  = F.softmax(scaled, dim=-1)116 117        # At final step use greedy (argmax) for cleaner output118        if step == sampling_steps:119            predicted = torch.argmax(probs, dim=-1)120        else:121            predicted = torch.multinomial(probs, 1).squeeze(-1)122 123        # Confidence scoring (on unscaled logits for reliable scores)124        true_probs  = F.softmax(response_logits, dim=-1)125        confidences = true_probs[torch.arange(num_masks), predicted]126 127        # Update running confidence (exponential moving average)128        # This smooths out noisy single-step confidence and gives a better129        # signal for which tokens should be revealed vs remasked130        running_confidence = 0.7 * running_confidence + 0.3 * confidences131 132        current = seq_tensor.squeeze(0).clone()133        for i, idx in enumerate(mask_indices):134            current[idx] = predicted[i]135 136        # Progressive remasking based on running confidence137        if step < sampling_steps:138            # Reveal more tokens as we progress139            target_revealed = int(num_masks * step / sampling_steps)140            num_remask = num_masks - target_revealed141            if num_remask > 0:142                # Re-mask the tokens with LOWEST running confidence143                _, low_idx = torch.topk(running_confidence, k=num_remask, largest=False)144                for li in low_idx:145                    current[mask_indices[li]] = mask_id146 147        seq_tensor = current.unsqueeze(0)148 149    # Decode final response150    response_ids = seq_tensor[0][response_start:response_end].tolist()151    return decode_response(response_ids, tokenizer)152 153 154if __name__ == "__main__":155    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")156    print(f"Inference on: {device}\n")157 158    tokenizer, vocab, id2word = load_tokenizer_full()159 160    # Must match train.py architecture exactly161    model = MaskedDiffusionModel(162        vocab_size=len(vocab),163        d_model=256, nhead=8, num_layers=6,164        max_seq_len=MAX_SEQ_LENGTH,165        dropout=0.0,  # No dropout at inference166    ).to(device)167 168    # Prefer EMA checkpoint (saves the smoothed weights)169    ckpt = "diffusion_model_best.pth" if os.path.exists("diffusion_model_best.pth") else "diffusion_model.pth"170    try:171        model.load_state_dict(torch.load(ckpt, map_location=device))172        print(f"Loaded: {ckpt} (EMA weights)\n")173    except Exception as e:174        print(f"Error loading checkpoint: {e}")175        print("Make sure to retrain with the new architecture first!")176        exit()177 178    test_prompts = [179        "hi",180        "how are you",181        "what is your name",182        "tell me a joke",183        "what do you do for fun",184        "i had a bad day",185        " Can we go now?",186    ]187 188    for prompt in test_prompts:189        response = generate_response(190            model, tokenizer, id2word,191            prompt=prompt,192            max_response_length=24,193            sampling_steps=40,194            temperature=0.5,195            top_k=15196        )197        print(f"User : {prompt}")198        print(f"Bot  : {response}")199        print("-" * 40)