Update app.py

app.py

@@ -3,7 +3,10 @@ import torch
 import spaces
 import gradio as gr
 
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
 
 
 # =====================================================
@@ -15,32 +18,33 @@ MODEL_ID = os.getenv(
     "WeiboAI/VibeThinker-3B"
 )
 
-
-DEVICE = (
-    "cuda"
-    if torch.cuda.is_available()
-    else "cpu"
-)
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
+print("=" * 60)
+print("X-RUDRA MODEL SPACE")
 print("MODEL:", MODEL_ID)
 print("DEVICE:", DEVICE)
+print("=" * 60)
 
 
 # =====================================================
-# LOAD
+# LOAD TOKENIZER
 # =====================================================
 
 print("Loading tokenizer...")
 
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
-    trust_remote_code=True,
+    trust_remote_code=True
 )
 
 
-print("Loading model...")
+# =====================================================
+# LOAD MODEL
+# =====================================================
 
+print("Loading model...")
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
@@ -56,12 +60,11 @@ model = AutoModelForCausalLM.from_pretrained(
 
 model.eval()
 
-
 print("MODEL READY")
 
 
 # =====================================================
-# GENERATE
+# GENERATION
 # =====================================================
 
 
@@ -69,26 +72,24 @@ print("MODEL READY")
 def generate(
     prompt,
     max_tokens,
-    temperature,
+    temperature
 ):
 
     if not prompt.strip():
-        return "Enter prompt"
+        return "Please enter a prompt."
 
 
-    # Normal tokenizer input
-
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
         truncation=True,
-        max_length=4096,
+        max_length=4096
     )
 
 
     inputs = {
-        k: v.to(model.device)
-        for k, v in inputs.items()
+        key: value.to(model.device)
+        for key, value in inputs.items()
     }
 
 
@@ -100,7 +101,7 @@ def generate(
 
     with torch.no_grad():
 
-        output = model.generate(
+        outputs = model.generate(
 
             **inputs,
 
@@ -118,7 +119,7 @@ def generate(
 
             top_k=50,
 
-            repetition_penalty=1.1,
+            repetition_penalty=1.15,
 
             no_repeat_ngram_size=3,
 
@@ -132,18 +133,18 @@ def generate(
         )
 
 
-    answer = tokenizer.decode(
-        output[0][input_length:],
-        skip_special_tokens=True,
+    result = tokenizer.decode(
+        outputs[0][input_length:],
+        skip_special_tokens=True
     )
 
 
-    return answer.strip()
+    return result.strip()
 
 
 
 # =====================================================
-# GRADIO CHAT
+# CHAT FUNCTION
 # =====================================================
 
 
@@ -151,24 +152,35 @@ def chat(
     message,
     history,
     max_tokens,
-    temperature,
+    temperature
 ):
 
-    response = generate(
+    answer = generate(
         message,
         max_tokens,
-        temperature,
+        temperature
     )
 
+
     history = history or []
 
+
     history.append(
-        (
-            message,
-            response,
-        )
+        {
+            "role": "user",
+            "content": message
+        }
     )
 
+
+    history.append(
+        {
+            "role": "assistant",
+            "content": answer
+        }
+    )
+
+
     return "", history
 
 
@@ -179,93 +191,102 @@ def chat(
 
 
 with gr.Blocks(
-    title="X-RUDRA MODEL"
+    title="X-RUDRA"
 ) as demo:
 
 
     gr.Markdown(
         f"""
-# ⚡ X-RUDRA MODEL
+# ⚡ X-RUDRA
+
+Model:
+
+`{MODEL_ID}`
 
-**{MODEL_ID}**
+Device:
 
-Device: `{DEVICE}`
+`{DEVICE}`
 """
     )
 
 
     chatbot = gr.Chatbot(
+        type="messages",
         height=600,
+        label="Chat"
     )
 
 
     with gr.Row():
 
         message = gr.Textbox(
-            placeholder="Ask something...",
-            scale=8,
+            placeholder="Ask anything...",
+            scale=8
         )
 
+
         send = gr.Button(
             "Send",
             variant="primary",
-            scale=1,
+            scale=1
         )
 
 
+
     with gr.Row():
 
         max_tokens = gr.Slider(
-            64,
-            2048,
+            minimum=64,
+            maximum=2048,
             value=512,
             step=64,
-            label="Max Tokens",
+            label="Max Tokens"
         )
 
 
         temperature = gr.Slider(
-            0.1,
-            1.2,
+            minimum=0.1,
+            maximum=1.2,
             value=0.7,
             step=0.1,
-            label="Temperature",
+            label="Temperature"
         )
 
 
+
     send.click(
-        chat,
+        fn=chat,
         inputs=[
             message,
             chatbot,
             max_tokens,
-            temperature,
+            temperature
         ],
         outputs=[
             message,
-            chatbot,
-        ],
+            chatbot
+        ]
     )
 
 
     message.submit(
-        chat,
+        fn=chat,
         inputs=[
             message,
             chatbot,
             max_tokens,
-            temperature,
+            temperature
         ],
         outputs=[
             message,
-            chatbot,
-        ],
+            chatbot
+        ]
     )
 
 
 
 # =====================================================
-# START
+# START SERVER
 # =====================================================
 
 if __name__ == "__main__":
@@ -273,4 +294,5 @@ if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
+        show_error=True
     )