Testing minimal interface

app.py

@@ -3,17 +3,16 @@ import torch
 import torch.nn.functional as F
 import time
 from tokenizers import Tokenizer
-import json
 from train2 import MaskedDiffusionModel
 
 MAX_SEQ_LENGTH = 64
 
 def load_model_and_tokenizer():
     device = torch.device("cpu")
+    # Using the exact same architecture as training
     tokenizer = Tokenizer.from_file("subword_tokenizer2.json")
     vocab = tokenizer.get_vocab()
-    id2word = {int(v): k for k, v in vocab.items()}
-
+    
     model = MaskedDiffusionModel(
         vocab_size=len(vocab),
         d_model=256,
@@ -23,17 +22,20 @@ def load_model_and_tokenizer():
         dropout=0.2
     ).to(device)
 
-    # Load weights (removing weights_only for better compatibility)
-    state_dict = torch.load("diffusion_model_between.pth", map_location=device)
-    model.load_state_dict(state_dict)
+    # Simplified load
+    try:
+        state_dict = torch.load("diffusion_model_between.pth", map_location=device)
+        model.load_state_dict(state_dict)
+    except Exception as e:
+        print(f"FAILED TO LOAD MODEL: {e}")
+        
     model.eval()
-    return model, tokenizer, id2word, device
+    return model, tokenizer, device
 
-model, tokenizer, id2word, device = load_model_and_tokenizer()
+model, tokenizer, device = load_model_and_tokenizer()
 
 def decode_with_masks(tensor, is_final=False):
     eos_id = tokenizer.token_to_id("[EOS]")
-    # Only truncate at EOS on the very final step, otherwise the text jumps around during refinement!
     if is_final:
         eos_indices = (tensor == eos_id).nonzero(as_tuple=True)[0]
         if len(eos_indices) > 0:
@@ -45,40 +47,27 @@ def decode_with_masks(tensor, is_final=False):
     filtered_ids = [tid for tid in tensor.tolist() if tid not in special_ids]
     if not filtered_ids: return ""
     
-    # Decode WITHOUT skipping special tokens so [MASK] is visible
     text = tokenizer.decode(filtered_ids, skip_special_tokens=False).strip()
-    
-    # Stylize the raw [MASK] token
     text = text.replace("[MASK]", "█")
     
-    # Text formatting correction
     for p in [".", ",", "?", "!", "'", ":"]:
         text = text.replace(f" {p}", p)
         
     return text.strip()
 
-def generate_step_by_step(history):
+def predict(message, history):
     try:
-        if not history:
-            return []
-        
-        prompt = history[-1][0]
-        print(f"Generating for prompt: {prompt}")
-        
-        steps = 20
-        temp = 0.2
+        steps = 15
+        temp = 0.3
         top_k = 10
 
         bos_id = tokenizer.token_to_id("[BOS]")
         eos_id = tokenizer.token_to_id("[EOS]")
         mask_id = tokenizer.token_to_id("[MASK]")
         pad_id = tokenizer.token_to_id("[PAD]")
-        unk_id = tokenizer.token_to_id("[UNK]")
-        special_ids = {bos_id, eos_id, mask_id, pad_id, unk_id}
-
-        formatted_prompt = f"user: {prompt.lower().strip()} bot:"
+        
+        formatted_prompt = f"user: {message.lower().strip()} bot:"
         input_ids = tokenizer.encode(formatted_prompt).ids
-
         max_resp = min(40, MAX_SEQ_LENGTH - len(input_ids) - 2)
 
         sequence = [bos_id] + input_ids + [mask_id] * max_resp + [eos_id]
@@ -91,292 +80,65 @@ def generate_step_by_step(history):
         num_masks = len(mask_indices)
 
         running_confidence = torch.zeros(num_masks, device=device)
-
-        def get_output(seq_tensor, is_final=False):
-            response_ids = seq_tensor[0][response_start:response_end]
-            return decode_with_masks(response_ids, is_final)
-
-        out_text = get_output(seq_tensor, is_final=False)
-        history[-1][1] = out_text
-        yield history
-
         current_seq = seq_tensor.squeeze(0).clone()
 
-        # Pre-calculate penalty indices once
-        vocab_size = model.fc_out.out_features
-        
+        output_text = ""
         for step in range(1, steps + 1):
             t_val = max(1.0 - step / steps, 0.05)
             t = torch.tensor([t_val], device=device)
-            pad_mask = (seq_tensor == pad_id)
             
             with torch.no_grad():
-                logits = model(seq_tensor, t, src_key_padding_mask=pad_mask)
-                
+                logits = model(seq_tensor, t, src_key_padding_mask=(seq_tensor == pad_id))
+            
             response_logits = logits[0, mask_indices]
             
-            # Optimized Vectorized Repetition Penalty
+            # Simple Repetition Penalty
             if step > 1:
-                prev_pred = current_seq[mask_indices]
-                # Get counts only for tokens present in context to avoid giant loops
-                unique_tokens, counts = torch.unique(prev_pred, return_counts=True)
+                unique_tokens, counts = torch.unique(current_seq[mask_indices], return_counts=True)
                 for i, tok_id in enumerate(unique_tokens):
                     t_id = tok_id.item()
-                    if t_id not in special_ids and counts[i] > 1:
-                        penalty = 5.0 ** (counts[i].item() - 1)
-                        response_logits[:, t_id] = torch.where(
-                            response_logits[:, t_id] < 0,
-                            response_logits[:, t_id] * penalty,
-                            response_logits[:, t_id] / penalty
-                        )
+                    if t_id not in [bos_id, eos_id, mask_id, pad_id] and counts[i] > 1:
+                        response_logits[:, t_id] /= 2.0
             
             if top_k > 0:
-                top_k_vals, _ = torch.topk(response_logits, top_k, dim=-1)
-                min_top_k = top_k_vals[:, -1].unsqueeze(-1)
-                response_logits = response_logits.masked_fill(response_logits < min_top_k, float('-inf'))
-                
-            probs = F.softmax(response_logits / max(temp, 1e-6), dim=-1)
+                v, _ = torch.topk(response_logits, top_k)
+                response_logits[response_logits < v[:, -1].unsqueeze(-1)] = -float('Inf')
+            
+            probs = F.softmax(response_logits / temp, dim=-1)
             
-            # Nan protection and sampling
-            if torch.isnan(probs).any():
-                print("WARNING: NaNs detected in probabilities, falling back to uniform")
-                probs = torch.ones_like(probs) / probs.size(-1)
-
             if step == steps:
                 predicted = torch.argmax(probs, dim=-1)
             else:
-                try:
-                    predicted = torch.multinomial(probs, 1).squeeze(-1)
-                except RuntimeError as e:
-                    print(f"Sampling error: {e}")
-                    predicted = torch.argmax(probs, dim=-1) # fallback to greedy
-                
-            true_probs = F.softmax(response_logits, dim=-1)
-            confidences = true_probs[torch.arange(num_masks), predicted]
+                predicted = torch.multinomial(probs, 1).squeeze(-1)
             
+            confidences = torch.gather(F.softmax(response_logits, dim=-1), 1, predicted.unsqueeze(-1)).squeeze(-1)
             running_confidence = 0.7 * running_confidence + 0.3 * confidences
             
             for i, idx in enumerate(mask_indices):
                 current_seq[idx] = predicted[i]
                 
             if step < steps:
-                target_revealed = int(num_masks * step / steps)
-                num_remask = num_masks - target_revealed
-                if num_remask > 0:
-                    _, low_idx = torch.topk(running_confidence, k=num_remask, largest=False)
+                target_reveal = int(num_masks * step / steps)
+                remask_count = num_masks - target_reveal
+                if remask_count > 0:
+                    _, low_idx = torch.topk(running_confidence, k=remask_count, largest=False)
                     for li in low_idx:
                         current_seq[mask_indices[li]] = mask_id
-                        
+            
             seq_tensor = current_seq.unsqueeze(0)
-            out_text = get_output(seq_tensor, is_final=(step == steps))
+            output_text = decode_with_masks(current_seq[response_start:response_end], is_final=(step == steps))
+            yield output_text
             
-            history[-1][1] = out_text
-            yield history
-            time.sleep(0.02)
     except Exception as e:
-        import traceback
-        error_msg = f"Inference Error: {str(e)}"
-        print(error_msg)
-        print(traceback.format_exc())
-        history[-1][1] = error_msg
-        yield history
-
-
-custom_css = """
-@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;800&family=JetBrains+Mono&display=swap');
-
-body, .gradio-container {
-    background-color: #030509 !important;
-    font-family: 'Inter', sans-serif !important;
-    color: #ffffff !important;
-    min-height: 100vh;
-    margin: 0;
-}
-
-@keyframes float-glow {
-    0% { transform: scale(1) translate(0, 0); opacity: 0.5; }
-    50% { transform: scale(1.2) translate(-20px, 20px); opacity: 0.8; }
-    100% { transform: scale(1) translate(0, 0); opacity: 0.5; }
-}
-
-.bg-mesh-1, .bg-mesh-2 {
-    position: fixed;
-    border-radius: 50%;
-    filter: blur(100px);
-    z-index: -1;
-    pointer-events: none;
-}
-.bg-mesh-1 {
-    width: 50vw; height: 50vw;
-    background: radial-gradient(circle, rgba(100, 150, 255, 0.4) 0%, transparent 60%);
-    top: -10vw; left: -10vw;
-    animation: float-glow 20s ease-in-out infinite;
-}
-.bg-mesh-2 {
-    width: 60vw; height: 60vw;
-    background: radial-gradient(circle, rgba(50, 100, 250, 0.3) 0%, transparent 60%);
-    bottom: -20vw; right: -20vw;
-    animation: float-glow 25s ease-in-out infinite reverse;
-}
-
-.hero-container {
-    padding: 6vh 5vw 2vh 5vw;
-    z-index: 10;
-}
-.hero-brand {
-    display: inline-flex;
-    align-items: center;
-    background: rgba(255, 255, 255, 0.1);
-    padding: 8px 16px;
-    border-radius: 20px;
-    font-size: 0.8rem;
-    font-weight: 600;
-    letter-spacing: 1px;
-    color: #ffffff;
-    margin-bottom: 2rem;
-    backdrop-filter: blur(10px);
-    border: 1px solid rgba(255, 255, 255, 0.15);
-}
-.hero-title {
-    font-size: clamp(3rem, 6vw, 6rem);
-    font-weight: 800;
-    line-height: 1.0;
-    letter-spacing: -0.04em;
-    color: #fff;
-    margin: 0 0 1rem 0;
-}
-.hero-description {
-    font-size: clamp(1rem, 1.2vw, 1.1rem);
-    color: #9ba3af;
-    font-weight: 400;
-    line-height: 1.6;
-    max-width: 900px;
-    margin-bottom: 2rem;
-    background: rgba(10, 15, 25, 0.4);
-    border-radius: 12px;
-    backdrop-filter: blur(10px);
-}
-.app-container {
-    padding: 0 5vw 10vh 5vw;
-    z-index: 10;
-}
-.chatbot-container {
-    border-radius: 24px !important;
-    background: rgba(15, 20, 25, 0.6) !important;
-    backdrop-filter: blur(25px) !important;
-    border: 1px solid rgba(255, 255, 255, 0.08) !important;
-    box-shadow: 0 20px 50px rgba(0,0,0,0.5), inset 0 1px 0 rgba(255,255,255,0.1) !important;
-}
-.message-wrap .user {
-    background: linear-gradient(135deg, rgba(255,255,255,0.1), rgba(255,255,255,0.05)) !important;
-    border: 1px solid rgba(255, 255, 255, 0.1) !important;
-    box-shadow: 0 4px 15px rgba(0,0,0,0.2) !important;
-}
-.message-wrap .bot {
-    background: transparent !important;
-    border: none !important;
-}
-
-/* Fancy Input Row */
-#input-container {
-    background: rgba(20, 25, 30, 0.7);
-    backdrop-filter: blur(20px);
-    border-radius: 24px;
-    border: 1px solid rgba(255, 255, 255, 0.15);
-    padding: 12px 12px 12px 24px;
-    margin-top: 20px;
-    box-shadow: 0 10px 30px rgba(0,0,0,0.3);
-    transition: all 0.3s ease;
-}
-#input-container:focus-within {
-    border-color: rgba(100, 200, 255, 0.5);
-    box-shadow: 0 0 30px rgba(100, 200, 255, 0.15);
-}
-#textbox textarea {
-    font-family: 'Inter', sans-serif !important;
-    font-size: 1.1rem !important;
-    color: white !important;
-    padding: 5px 0 0 0 !important;
-}
-#textbox textarea::placeholder {
-    color: rgba(255,255,255,0.4) !important;
-}
-#submit-btn {
-    color: #000 !important;
-    border-radius: 16px !important;
-    font-weight: 700 !important;
-    border: none !important;
-    box-shadow: 0 5px 15px rgba(0, 242, 254, 0.3) !important;
-    transition: all 0.3s ease !important;
-}
-
-/* Examples Styling */
-.gr-examples {
-    margin-top: 25px !important;
-}
-.gr-examples h3, .gr-examples .label {
-    color: rgba(255,255,255,0.6) !important;
-    font-weight: 500 !important;
-    font-size: 0.9rem !important;
-}
-.gr-examples .gallery-item {
-    background: rgba(255,255,255,0.05) !important;
-    border: 1px solid rgba(255,255,255,0.1) !important;
-    color: white !important;
-    border-radius: 12px !important;
-    transition: background 0.3s !important;
-}
-.gr-examples .gallery-item:hover {
-    background: rgba(255,255,255,0.15) !important;
-}
-"""
-
-with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:
-    gr.HTML("""
-    <div class="bg-mesh-1"></div>
-    <div class="bg-mesh-2"></div>
-    """)
-    
-    with gr.Column(elem_classes="hero-container"):
-        gr.HTML("""
-        
-        <h1 class="hero-title">PersonaChat MDLM</h1>
-        <div class="hero-description">
-            <strong>Architecture:</strong> 17 Million Parameter Masked Discrete Diffusion Language Model<br/><br/>
-            Unlike traditional autoregressive models (like ChatGPT) that sequentially guess words one-by-one, this model harnesses <strong>Parallel Denoising Generation</strong>. It maps out the entire sequence's layout instantly, then uses progressive <em>Timestep Adaptive Layer Normalization (AdaLN)</em> and <em>EMA Confidence Trajectory</em> to morph static noise into coherent sentences.<br/><br/>
-            <strong>Trained on:</strong> <code>bavard/personachat_truecased</code> dataset.
-        </div>
-        <div class="hero-brand">
-            DIFFUSION MODEL
-        </div>
-        """)
-
-    with gr.Column(elem_classes="app-container"):
-        chatbot = gr.Chatbot(elem_classes="chatbot-container", height=450, show_label=False)
-        
-        with gr.Row(elem_id="input-container"):
-            msg = gr.Textbox(placeholder="Ask me something about myself...", show_label=False, elem_id="textbox", scale=8, container=False)
-            submit = gr.Button("Query", elem_id="submit-btn", scale=1)
-            
-        gr.Examples(
-            examples=[
-                "Hi, how are you doing today?",
-                "How are you doing?",
-                "Do you have any pets?",
-                "What kind of music do you like?"
-            ],
-            inputs=msg
-        )
-
-    def user(user_message, history):
-        return "", history + [[user_message, None]]
-
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        generate_step_by_step, chatbot, chatbot
-    )
-    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        generate_step_by_step, chatbot, chatbot
-    )
+        yield f"Error: {str(e)}"
+
+demo = gr.ChatInterface(
+    predict,
+    title="PersonaChat MDLM",
+    description="Minimal diagnostic interface. If this works, the problem was the CSS/UI complexity.",
+    examples=["Hi, how are you?", "What are your hobbies?", "Do you have any pets?"],
+    theme="soft"
+)
 
 if __name__ == "__main__":
     demo.launch()