Update app.py

app.py

@@ -3,21 +3,19 @@ import torch
 import spaces
 import gradio as gr
 
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
-# ============================================================
+# =====================================================
 # CONFIG
-# ============================================================
+# =====================================================
 
 MODEL_ID = os.getenv(
     "MODEL_ID",
     "WeiboAI/VibeThinker-3B"
 )
 
+
 DEVICE = (
     "cuda"
     if torch.cuda.is_available()
@@ -25,16 +23,13 @@ DEVICE = (
 )
 
 
-print("=" * 60)
-print("X-RUDRA MODEL SPACE")
 print("MODEL:", MODEL_ID)
 print("DEVICE:", DEVICE)
-print("=" * 60)
 
 
-# ============================================================
-# LOAD MODEL
-# ============================================================
+# =====================================================
+# LOAD
+# =====================================================
 
 print("Loading tokenizer...")
 
@@ -49,12 +44,12 @@ print("Loading model...")
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
+    device_map="auto",
     torch_dtype=(
         torch.float16
         if DEVICE == "cuda"
         else torch.float32
     ),
-    device_map="auto",
     trust_remote_code=True,
 )
 
@@ -65,9 +60,9 @@ model.eval()
 print("MODEL READY")
 
 
-# ============================================================
-# GENERATION
-# ============================================================
+# =====================================================
+# GENERATE
+# =====================================================
 
 
 @spaces.GPU
@@ -78,15 +73,16 @@ def generate(
 ):
 
     if not prompt.strip():
+        return "Enter prompt"
 
-        return "Enter prompt."
 
+    # Normal tokenizer input
 
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
-        padding=True,
         truncation=True,
+        max_length=4096,
     )
 
 
@@ -96,7 +92,6 @@ def generate(
     }
 
 
-    # FIXED
     input_length = (
         inputs["input_ids"]
         .shape[-1]
@@ -105,7 +100,7 @@ def generate(
 
     with torch.no_grad():
 
-        outputs = model.generate(
+        output = model.generate(
 
             **inputs,
 
@@ -119,6 +114,18 @@ def generate(
 
             do_sample=True,
 
+            top_p=0.95,
+
+            top_k=50,
+
+            repetition_penalty=1.1,
+
+            no_repeat_ngram_size=3,
+
+            eos_token_id=(
+                tokenizer.eos_token_id
+            ),
+
             pad_token_id=(
                 tokenizer.eos_token_id
             ),
@@ -126,7 +133,7 @@ def generate(
 
 
     answer = tokenizer.decode(
-        outputs[0][input_length:],
+        output[0][input_length:],
         skip_special_tokens=True,
     )
 
@@ -135,38 +142,77 @@ def generate(
 
 
 
-# ============================================================
+# =====================================================
+# GRADIO CHAT
+# =====================================================
+
+
+def chat(
+    message,
+    history,
+    max_tokens,
+    temperature,
+):
+
+    response = generate(
+        message,
+        max_tokens,
+        temperature,
+    )
+
+    history = history or []
+
+    history.append(
+        (
+            message,
+            response,
+        )
+    )
+
+    return "", history
+
+
+
+# =====================================================
 # UI
-# ============================================================
+# =====================================================
 
 
 with gr.Blocks(
-    title="X-RUDRA M2"
+    title="X-RUDRA MODEL"
 ) as demo:
 
 
     gr.Markdown(
         f"""
-# ⚡ X-RUDRA M2
+# ⚡ X-RUDRA MODEL
 
-Model:
+**{MODEL_ID}**
 
-`{MODEL_ID}`
-
-Device:
-
-`{DEVICE}`
+Device: `{DEVICE}`
 """
     )
 
 
-    prompt = gr.Textbox(
-        label="Prompt",
-        lines=6,
-        placeholder="Ask something..."
+    chatbot = gr.Chatbot(
+        height=600,
     )
 
 
+    with gr.Row():
+
+        message = gr.Textbox(
+            placeholder="Ask something...",
+            scale=8,
+        )
+
+        send = gr.Button(
+            "Send",
+            variant="primary",
+            scale=1,
+        )
+
+
     with gr.Row():
 
         max_tokens = gr.Slider(
@@ -174,48 +220,56 @@ Device:
             2048,
             value=512,
             step=64,
-            label="Max Tokens"
+            label="Max Tokens",
         )
 
 
         temperature = gr.Slider(
             0.1,
-            1.5,
+            1.2,
             value=0.7,
             step=0.1,
-            label="Temperature"
+            label="Temperature",
         )
 
 
-    btn = gr.Button(
-        "Generate",
-        variant="primary"
-    )
-
-
-    output = gr.Markdown(
-        label="Response"
+    send.click(
+        chat,
+        inputs=[
+            message,
+            chatbot,
+            max_tokens,
+            temperature,
+        ],
+        outputs=[
+            message,
+            chatbot,
+        ],
     )
 
 
-    btn.click(
-        generate,
+    message.submit(
+        chat,
         inputs=[
-            prompt,
+            message,
+            chatbot,
             max_tokens,
             temperature,
         ],
-        outputs=output,
+        outputs=[
+            message,
+            chatbot,
+        ],
     )
 
 
-# ============================================================
+
+# =====================================================
 # START
-# ============================================================
+# =====================================================
 
 if __name__ == "__main__":
 
-
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,