Update app.py

app.py

@@ -1,54 +1,136 @@
 import os
 import torch
+import spaces
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# =====================================================
+# ============================================================
 # CONFIG
-# =====================================================
+# ============================================================
 
 MODEL_ID = os.getenv("MODEL_ID", "WeiboAI/VibeThinker-3B")
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 print("=" * 60)
-print("X-RUDRA MODEL SPACE")
+print("X-RUDRA M2 (CHAT)")
 print("MODEL:", MODEL_ID)
 print("DEVICE:", DEVICE)
 print("=" * 60)
 
-# =====================================================
-# LOAD TOKENIZER & MODEL
-# =====================================================
+# ============================================================
+# LOAD MODEL
+# ============================================================
 
 print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
+# Ensure pad_token is set (some tokenizers don't have one)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
 print("Loading model...")
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
+    dtype=torch.float16 if DEVICE == "cuda" else torch.float32,   # FIXED: use dtype
     device_map="auto",
-    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
     trust_remote_code=True,
 )
 model.eval()
 print("MODEL READY")
 
-# =====================================================
-# GENERATION (unchanged)
-# =====================================================
+# ============================================================
+# HELPERS: CONVERSATION FORMATTING
+# ============================================================
+
+def format_conversation(history):
+    """
+    Convert history (list of tuples [(user, assistant), ...]) into a single prompt.
+    Uses the tokenizer's chat template if available, else a simple User/Assistant format.
+    """
+    # Try to use the tokenizer's built-in chat template (if any)
+    if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
+        # Convert history to list of dicts
+        messages = []
+        for user_msg, assistant_msg in history:
+            messages.append({"role": "user", "content": user_msg})
+            if assistant_msg:
+                messages.append({"role": "assistant", "content": assistant_msg})
+        # Add the final user message placeholder? Actually we'll append the new user message separately.
+        # But we are calling this before adding the new user message? We'll design it to include the new user message.
+        # Better to handle in generate_response: we'll pass the full history including the new user message.
+        # So we'll restructure: generate_response will build the messages list and then call apply_chat_template.
+        # We'll move this logic into generate_response directly.
+        # So this function will only be used as fallback.
+        pass
+
+    # Fallback: simple User/Assistant format
+    prompt = ""
+    for user_msg, assistant_msg in history:
+        prompt += f"User: {user_msg}\n"
+        if assistant_msg:
+            prompt += f"Assistant: {assistant_msg}\n"
+    # We will append the new user message outside this function
+    return prompt
+
+# ============================================================
+# GENERATION (with chat history)
+# ============================================================
 
-def generate(prompt, max_tokens, temperature):
-    if not prompt.strip():
-        return "Please enter a prompt."
+@spaces.GPU
+def generate_response(message, history, max_tokens, temperature):
+    """
+    Takes the current message and history (list of tuples), generates a response,
+    and returns (new_message, updated_history) – new_message is always '' to clear the input.
+    """
+    history = history or []
 
+    # Build the full conversation prompt
+    # If the tokenizer has a chat template, use it
+    if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
+        # Convert history (tuples) to messages list
+        messages = []
+        for user_msg, assistant_msg in history:
+            messages.append({"role": "user", "content": user_msg})
+            if assistant_msg:
+                messages.append({"role": "assistant", "content": assistant_msg})
+        # Add the new user message
+        messages.append({"role": "user", "content": message})
+        # Apply the template
+        try:
+            full_prompt = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+        except Exception as e:
+            print("Chat template failed, falling back to manual format:", e)
+            full_prompt = None
+    else:
+        full_prompt = None
+
+    # If template failed or not available, use manual format
+    if full_prompt is None:
+        # Build from history + new message
+        prompt = ""
+        for user_msg, assistant_msg in history:
+            prompt += f"User: {user_msg}\n"
+            if assistant_msg:
+                prompt += f"Assistant: {assistant_msg}\n"
+        # Add the new user message
+        prompt += f"User: {message}\nAssistant:"
+        full_prompt = prompt
+
+    # Tokenize
     inputs = tokenizer(
-        prompt,
+        full_prompt,
         return_tensors="pt",
         truncation=True,
-        max_length=4096
+        max_length=4096,          # adjust if needed
+        padding=True,
     )
     inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    input_length = inputs["input_ids"].shape[-1]
+
+    input_len = inputs["input_ids"].shape[-1]
 
     with torch.no_grad():
         outputs = model.generate(
@@ -60,76 +142,58 @@ def generate(prompt, max_tokens, temperature):
             top_k=50,
             repetition_penalty=1.15,
             no_repeat_ngram_size=3,
+            pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.eos_token_id,
         )
 
-    result = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
-    return result.strip()
-
-# =====================================================
-# CONVERSATION FORMATTER
-# =====================================================
-
-def format_conversation(history):
-    """Convert history (list of dicts) into a single prompt string."""
-    prompt = ""
-    for turn in history:
-        if turn["role"] == "user":
-            prompt += f"User: {turn['content']}\n"
-        elif turn["role"] == "assistant":
-            prompt += f"Assistant: {turn['content']}\n"
-    # Append the final assistant prompt to continue the conversation
-    prompt += "Assistant:"
-    return prompt
-
-# =====================================================
-# CHAT FUNCTION (FIXED)
-# =====================================================
-
-def chat(message, history, max_tokens, temperature):
-    history = history or []
-
-    # Append the new user message to history
-    history.append({"role": "user", "content": message})
+    # Decode only new tokens
+    new_tokens = outputs[0][input_len:]
+    answer = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
-    # Build the full conversation prompt
-    full_prompt = format_conversation(history)
-
-    # Generate assistant response from the full context
-    answer = generate(full_prompt, max_tokens, temperature)
+    # Update history: add (user, assistant) tuple
+    history.append((message, answer))
 
-    # Append assistant response to history
-    history.append({"role": "assistant", "content": answer})
+    return "", history   # clear input, return updated history
 
-    return "", history
+# ============================================================
+# UI (Chat Interface – using tuples for Gradio 3.x)
+# ============================================================
 
-# =====================================================
-# UI (unchanged)
-# =====================================================
+with gr.Blocks(title="X-RUDRA M2") as demo:
+    gr.Markdown(
+        f"""
+# ⚡ X-RUDRA M2 – Chat
+**Model:** `{MODEL_ID}`  
+**Device:** `{DEVICE}`
+"""
+    )
 
-with gr.Blocks(title="X-RUDRA") as demo:
-    gr.Markdown(f"""
-# ⚡ X-RUDRA
-Model: `{MODEL_ID}`  
-Device: `{DEVICE}`
-""")
+    # Chatbot component – no 'type' argument
+    chatbot = gr.Chatbot(height=600, label="Conversation")
 
-    chatbot = gr.Chatbot(type="messages", height=600, label="Chat")
     with gr.Row():
-        message = gr.Textbox(placeholder="Ask anything...", scale=8)
+        msg = gr.Textbox(placeholder="Ask anything...", scale=8)
         send = gr.Button("Send", variant="primary", scale=1)
 
     with gr.Row():
-        max_tokens = gr.Slider(minimum=64, maximum=2048, value=512, step=64, label="Max Tokens")
-        temperature = gr.Slider(minimum=0.1, maximum=1.2, value=0.7, step=0.1, label="Temperature")
-
-    send.click(fn=chat, inputs=[message, chatbot, max_tokens, temperature], outputs=[message, chatbot])
-    message.submit(fn=chat, inputs=[message, chatbot, max_tokens, temperature], outputs=[message, chatbot])
+        max_tokens = gr.Slider(64, 2048, value=512, step=64, label="Max Tokens")
+        temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
+
+    # When Send is clicked or Enter is pressed
+    send.click(
+        fn=generate_response,
+        inputs=[msg, chatbot, max_tokens, temperature],
+        outputs=[msg, chatbot]
+    )
+    msg.submit(
+        fn=generate_response,
+        inputs=[msg, chatbot, max_tokens, temperature],
+        outputs=[msg, chatbot]
+    )
 
-# =====================================================
-# START SERVER
-# =====================================================
+# ============================================================
+# START
+# ============================================================
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860)