ritikraj2425/Discrete-Diffusion-Text-Demo

1import gradio as gr2import torch3import torch.nn.functional as F4import time5from tokenizers import Tokenizer6from train2 import MaskedDiffusionModel7 8MAX_SEQ_LENGTH = 649 10def load_model_and_tokenizer():11    device = torch.device("cpu")12    tokenizer = Tokenizer.from_file("subword_tokenizer2.json")13    vocab = tokenizer.get_vocab()14    15    model = MaskedDiffusionModel(16        vocab_size=len(vocab),17        d_model=256,18        nhead=8,19        num_layers=6,20        max_seq_len=MAX_SEQ_LENGTH,21        dropout=0.222    ).to(device)23 24    try:25        state_dict = torch.load("diffusion_model_between.pth", map_location=device)26        model.load_state_dict(state_dict)27    except Exception as e:28        print(f"FAILED TO LOAD MODEL: {e}")29        30    model.eval()31    return model, tokenizer, device32 33model, tokenizer, device = load_model_and_tokenizer()34 35def decode_with_masks(tensor, is_final=False):36    eos_id = tokenizer.token_to_id("[EOS]")37    if is_final:38        eos_indices = (tensor == eos_id).nonzero(as_tuple=True)[0]39        if len(eos_indices) > 0:40            tensor = tensor[:eos_indices[0]]41 42    special_ids = {tokenizer.token_to_id("[PAD]"), tokenizer.token_to_id("[BOS]"),43                   tokenizer.token_to_id("[EOS]"), tokenizer.token_to_id("[UNK]")}44    45    filtered_ids = [tid for tid in tensor.tolist() if tid not in special_ids]46    if not filtered_ids: return ""47    48    text = tokenizer.decode(filtered_ids, skip_special_tokens=False).strip()49    text = text.replace("[MASK]", "█")50    51    for p in [".", ",", "?", "!", "'", ":"]:52        text = text.replace(f" {p}", p)53        54    return text.strip()55 56def predict(message, history):57    try:58        steps = 1559        temp = 0.360        top_k = 1061 62        bos_id = tokenizer.token_to_id("[BOS]")63        eos_id = tokenizer.token_to_id("[EOS]")64        mask_id = tokenizer.token_to_id("[MASK]")65        pad_id = tokenizer.token_to_id("[PAD]")66        67        formatted_prompt = f"user: {message.lower().strip()} bot:"68        input_ids = tokenizer.encode(formatted_prompt).ids69        max_resp = min(40, MAX_SEQ_LENGTH - len(input_ids) - 2)70 71        sequence = [bos_id] + input_ids + [mask_id] * max_resp + [eos_id]72        sequence += [pad_id] * (MAX_SEQ_LENGTH - len(sequence))73        seq_tensor = torch.tensor([sequence], dtype=torch.long, device=device)74 75        response_start = 1 + len(input_ids)76        response_end = response_start + max_resp77        mask_indices = list(range(response_start, response_end))78        num_masks = len(mask_indices)79 80        running_confidence = torch.zeros(num_masks, device=device)81        current_seq = seq_tensor.squeeze(0).clone()82 83        output_text = ""84        for step in range(1, steps + 1):85            t_val = max(1.0 - step / steps, 0.05)86            t = torch.tensor([t_val], device=device)87            88            with torch.no_grad():89                logits = model(seq_tensor, t, src_key_padding_mask=(seq_tensor == pad_id))90            91            response_logits = logits[0, mask_indices]92            93            if step > 1:94                unique_tokens, counts = torch.unique(current_seq[mask_indices], return_counts=True)95                for i, tok_id in enumerate(unique_tokens):96                    t_id = tok_id.item()97                    if t_id not in [bos_id, eos_id, mask_id, pad_id] and counts[i] > 1:98                        # Logit subtraction entirely prevents structural duplication loops natively99                        response_logits[:, t_id] -= 10.0 * (counts[i].item() - 1)100            101            if top_k > 0:102                v, _ = torch.topk(response_logits, top_k)103                response_logits[response_logits < v[:, -1].unsqueeze(-1)] = -float('Inf')104            105            probs = F.softmax(response_logits / temp, dim=-1)106            107            if step == steps:108                predicted = torch.argmax(probs, dim=-1)109            else:110                predicted = torch.multinomial(probs, 1).squeeze(-1)111            112            confidences = torch.gather(F.softmax(response_logits, dim=-1), 1, predicted.unsqueeze(-1)).squeeze(-1)113            running_confidence = 0.7 * running_confidence + 0.3 * confidences114            115            for i, idx in enumerate(mask_indices):116                current_seq[idx] = predicted[i]117                118            if step < steps:119                target_reveal = int(num_masks * step / steps)120                remask_count = num_masks - target_reveal121                if remask_count > 0:122                    _, low_idx = torch.topk(running_confidence, k=remask_count, largest=False)123                    for li in low_idx:124                        current_seq[mask_indices[li]] = mask_id125            126            seq_tensor = current_seq.unsqueeze(0)127            output_text = decode_with_masks(current_seq[response_start:response_end], is_final=(step == steps))128            yield output_text129            130    except Exception as e:131        yield f"Error: {str(e)}"132 133custom_css = """134@import url('https://fonts.googleapis.com/css2?family=Fira+Code:wght@400;600&display=swap');135 136body, .gradio-container {137    background-color: #1e1e1e !important;138    font-family: 'Fira Code', monospace !important;139    color: #d4d4d4 !important;140}141 142.hero-container {143    padding: 2rem 5vw;144    border-bottom: 2px solid #333;145    background: #191919;146}147 148.hero-brand {149    color: #569cd6;150    font-size: 1rem;151    margin-bottom: 0.5rem;152}153 154.hero-brand::before { content: "<"; color: #808080; }155.hero-brand::after { content: "/>"; color: #808080; }156 157.hero-title {158    font-size: 2.5rem;159    color: #ce9178;160    margin: 0 0 1rem 0;161    font-weight: 600;162}163 164.hero-description {165    color: #6a9955;166    line-height: 1.6;167    font-size: 1rem;168    background: transparent;169    padding: 0;170    border: none;171}172 173.hero-description strong {174    color: #c586c0;175}176 177.hero-description code {178    color: #dcdcaa;179    background: #2d2d2d;180    padding: 2px 6px;181    border-radius: 3px;182}183 184.app-container {185    padding: 0 5vw 5vh 5vw;186}187 188/* Customizing ChatInterface objects */189.bubble-wrap {190    font-family: 'Fira Code', monospace !important;191}192 193.message-wrap .user {194    background: #252526 !important;195    border: 1px solid #3c3c3c !important;196    color: #9cdcfe !important;197}198 199.message-wrap .bot {200    background: transparent !important;201    border: none !important;202    color: #d4d4d4 !important;203}204"""205 206with gr.Blocks(css=custom_css, fill_height=True) as demo:207    with gr.Column(elem_classes="hero-container"):208        gr.HTML("""209        <div class="hero-brand">210            persona-chat-mdlm211        </div>212        <h1 class="hero-title">PersonaChat MDLM</h1>213        <div class="hero-description">214            /*<br/>215            &nbsp;* <strong>Architecture:</strong> 17 Million Parameter Masked Discrete Diffusion Language Model<br/><br/>216            &nbsp;* Unlike traditional autoregressive models that guess words strictly left-to-right, this model employs <strong>Parallel Denoising Generation</strong>.<br/>217            &nbsp;* It maps out the structural sequence space instantly and iteratively normalizes masks into tokens.<br/><br/>218            &nbsp;* <strong>Speed Paradigm:</strong> True <code>O(1)</code> scaling factor. Because generation relies on parallel iterations,<br/>219            &nbsp;* computing a 10-token array demands the exact same temporal footprint as computing a 100-token array.<br/><br/>220            &nbsp;* <strong>Dataset Pipeline:</strong> <code>bavard/personachat_truecased</code><br/>221            &nbsp;*/222        </div>223        """)224 225    with gr.Column(elem_classes="app-container"):226        gr.ChatInterface(227            predict,228            examples=["Hi, how are you doing today?", "How are you doing?", "Do you have any pets?", "What kind of music do you like?"]229        )230 231if __name__ == "__main__":232    demo.launch()