corrected gradio output

app.py

@@ -1,104 +1,113 @@
 import gradio as gr
 import torch
 import torch.nn.functional as F
+import time
 from tokenizers import Tokenizer
 from train import MaskedDiffusionModel
 
-# --- Setup & Loading ---
+
 def load_model_and_tokenizer():
     device = torch.device("cpu")
+
     tokenizer = Tokenizer.from_file("subword_tokenizer.json")
     vocab = tokenizer.get_vocab()
     id2word = {int(v): k for k, v in vocab.items()}
-    
-    # Initialize with your last known settings
-    model = MaskedDiffusionModel(vocab_size=len(vocab), d_model=768, nhead=12, num_layers=12, max_seq_len=128).to(device)
+
+    model = MaskedDiffusionModel(
+        vocab_size=len(vocab),
+        d_model=768,
+        nhead=12,
+        num_layers=12,
+        max_seq_len=128
+    ).to(device)
+
     model.load_state_dict(torch.load("diffusion_model.pth", map_location=device))
     model.eval()
+
     return model, tokenizer, id2word, device
 
+
 model, tokenizer, id2word, device = load_model_and_tokenizer()
 
-def decode_tensor(tensor, id2word):
-    special_tokens = ["[PAD]", "[BOS]", "[EOS]", "[MASK]", "[UNK]"]
-    
+
+def decode_with_masks(tensor):
     words = []
     for token_id in tensor.tolist():
         word = id2word.get(token_id, "[UNK]")
-        if word not in special_tokens:
-            words.append(word)
-    
+        words.append(word)
     return " ".join(words).replace(" ##", "").replace("##", "")
 
-# --- Inference Generator ---
+
 def generate_step_by_step(prompt):
     max_len = 15
     steps = 15
-    temp = 0.1
-    
+    temp = 0.7
+
     bos_id = tokenizer.token_to_id("[BOS]")
     eos_id = tokenizer.token_to_id("[EOS]")
     mask_id = tokenizer.token_to_id("[MASK]")
-    
+
     formatted_prompt = f"user: {prompt} bot:"
     input_ids = tokenizer.encode(formatted_prompt).ids
-    
+
     sequence = [bos_id] + input_ids + [mask_id] * max_len + [eos_id]
     seq_tensor = torch.tensor([sequence], dtype=torch.long, device=device)
-    
+
     mask_indices = (seq_tensor == mask_id).squeeze(0).nonzero(as_tuple=True)[0]
-    
 
-    current_text = decode_tensor(seq_tensor[0], id2word)
-    yield f"**Step {step:02d}:** {current_text}"
+    def get_output(seq_tensor):
+        response_ids = seq_tensor[0][1 + len(input_ids):-1]
+        return decode_with_masks(response_ids)
+
+    # initial state
+    current_output = get_output(seq_tensor)
+    yield current_output
 
+    # diffusion loop
     for step in range(1, steps + 1):
         with torch.no_grad():
             logits = model(seq_tensor)
-        
+
         probs = F.softmax(logits / max(temp, 1e-6), dim=-1)
-        predicted_ids = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(probs.shape[:-1]).squeeze(0)
-        
+
+        predicted_ids = torch.multinomial(
+            probs.view(-1, probs.size(-1)), 1
+        ).view(probs.shape[:-1]).squeeze(0)
+
         true_probs = F.softmax(logits, dim=-1).squeeze(0)
-        confidences = torch.gather(true_probs, 1, predicted_ids.unsqueeze(1)).squeeze(1)
-        
+        confidences = torch.gather(
+            true_probs, 1, predicted_ids.unsqueeze(1)
+        ).squeeze(1)
+
         target_unmasked = int(len(mask_indices) * (step / steps))
         current_seq = seq_tensor.squeeze(0).clone()
-        
+
         for idx in mask_indices:
             current_seq[idx] = predicted_ids[idx]
-            
+
         if step < steps:
             gen_conf = confidences[mask_indices]
             num_remask = len(mask_indices) - target_unmasked
+
             if num_remask > 0:
                 _, low_conf_idx = torch.topk(gen_conf, k=num_remask, largest=False)
                 for idx in mask_indices[low_conf_idx]:
                     current_seq[idx] = mask_id
-        
-        seq_tensor = current_seq.unsqueeze(0)
-        output_log += f"`Step {step:02d}: {decode_tensor(seq_tensor[0], id2word)}`\n"
-        yield output_log
 
-    response_ids = seq_tensor[0][1 + len(input_ids) : -1]
-    final_text = decode_tensor(response_ids, id2word)
-    output_log += f"\n**Final Output:** {final_text}"
-    yield output_log
+        seq_tensor = current_seq.unsqueeze(0)
 
-# --- Gradio UI ---
-description_text = """
-### Experimental Discrete Diffusion Language Model (Proof of Concept)
-**Note to Reviewers:** This is an early-stage checkpoint of a 90M parameter Transformer trained from scratch using a Masked Discrete Diffusion objective. Due to compute access ending before convergence (Loss ~ 5.0), the final English generation is fragmented. 
+        current_output = get_output(seq_tensor)
+        yield current_output
 
-However, this demo successfully showcases the **Iterative Denoising Algorithm**. Watch how the model transitions from pure `[MASK]` tokens to predicted vocabulary over 15 discrete diffusion steps!
-"""
+        time.sleep(0.2)
+        
 
 demo = gr.Interface(
     fn=generate_step_by_step,
-    inputs=gr.Textbox(label="Enter a short prompt (e.g., 'hi')"),
-    outputs=gr.Markdown(label="Diffusion Denoising Steps"),
-    title="Discrete Diffusion NLP Demo",
-    description=description_text
+    inputs=gr.Textbox(label="Enter prompt"),
+    outputs=gr.Textbox(label="Generated Text (Converging)"),
+    title="Diffusion Model - Single Line Convergence",
+    description="Watch one sentence evolve from [MASK] tokens into final output."
 )
 
 if __name__ == "__main__":