Update app.py

app.py

@@ -1,9 +1,7 @@
 import os
-import time
-
 import torch
-import gradio as gr
 import spaces
+import gradio as gr
 
 from transformers import (
     AutoTokenizer,
@@ -17,10 +15,9 @@ from transformers import (
 
 MODEL_ID = os.getenv(
     "MODEL_ID",
-    "WeiboAI/VibeThinker-3B",
+    "WeiboAI/VibeThinker-3B"
 )
 
-
 DEVICE = (
     "cuda"
     if torch.cuda.is_available()
@@ -29,83 +26,43 @@ DEVICE = (
 
 
 print("=" * 60)
-print("X-RUDRA MODEL A")
+print("X-RUDRA MODEL SPACE")
 print("MODEL:", MODEL_ID)
 print("DEVICE:", DEVICE)
 print("=" * 60)
 
 
-
 # ============================================================
-# LAZY MODEL STORAGE
+# LOAD MODEL
 # ============================================================
 
-tokenizer = None
-model = None
-
-
-
-def load_model():
-
-    global tokenizer
-    global model
-
-
-    if model is not None:
-        return
-
-
-    print("Loading tokenizer...")
-
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_ID,
-        trust_remote_code=True,
-    )
-
-
-    model_kwargs = {
-        "trust_remote_code": True,
-    }
-
-
-    if torch.cuda.is_available():
+print("Loading tokenizer...")
 
-        model_kwargs["device_map"] = "auto"
-
-        model_kwargs["torch_dtype"] = (
-            torch.bfloat16
-        )
-
-    else:
-
-        model_kwargs["torch_dtype"] = (
-            torch.float32
-        )
-
-
-
-    print("Loading model...")
-
-
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        **model_kwargs,
-    )
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+)
 
 
-    model.eval()
+print("Loading model...")
 
 
-    if tokenizer.pad_token_id is None:
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=(
+        torch.float16
+        if DEVICE == "cuda"
+        else torch.float32
+    ),
+    device_map="auto",
+    trust_remote_code=True,
+)
 
-        tokenizer.pad_token = (
-            tokenizer.eos_token
-        )
 
+model.eval()
 
-    print("MODEL READY")
 
+print("MODEL READY")
 
 
 # ============================================================
@@ -113,329 +70,153 @@ def load_model():
 # ============================================================
 
 
-@spaces.GPU(
-    duration=120
-)
+@spaces.GPU
 def generate(
     prompt,
-    max_new_tokens,
+    max_tokens,
     temperature,
-    top_p,
 ):
 
+    if not prompt.strip():
 
-    if not prompt or not prompt.strip():
-
-        return (
-            "ERROR: prompt is empty."
-        )
-
-
-    load_model()
-
+        return "Enter prompt."
 
 
-    messages = [
-        {
-            "role": "user",
-            "content": prompt.strip(),
-        }
-    ]
-
-
-
-    if hasattr(
-        tokenizer,
-        "apply_chat_template"
-    ):
-
-
-        inputs = tokenizer.apply_chat_template(
-
-            messages,
-
-            tokenize=True,
-
-            add_generation_prompt=True,
-
-            return_tensors="pt",
-        )
-
-
-    else:
-
-
-        inputs = tokenizer(
-            prompt,
-            return_tensors="pt",
-        ).input_ids
-
-
-
-
-    if torch.cuda.is_available():
-
-        inputs = inputs.to(
-            model.device
-        )
-
-    else:
-
-        inputs = inputs.to(
-            DEVICE
-        )
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+    )
 
 
+    inputs = {
+        k: v.to(model.device)
+        for k, v in inputs.items()
+    }
 
-    input_length = inputs.shape[-1]
 
+    # FIXED
+    input_length = (
+        inputs["input_ids"]
+        .shape[-1]
+    )
 
 
-    with torch.inference_mode():
+    with torch.no_grad():
 
-        output = model.generate(
+        outputs = model.generate(
 
-            inputs,
+            **inputs,
 
             max_new_tokens=int(
-                max_new_tokens
+                max_tokens
             ),
 
             temperature=float(
                 temperature
             ),
 
-            top_p=float(
-                top_p
-            ),
-
-
-            do_sample=(
-                float(temperature) > 0
-            ),
-
+            do_sample=True,
 
             pad_token_id=(
-                tokenizer.pad_token_id
-            ),
-
-
-            eos_token_id=(
                 tokenizer.eos_token_id
             ),
         )
 
 
-
-    generated = output[0][
-        input_length:
-    ]
-
-
-
-    text = tokenizer.decode(
-
-        generated,
-
+    answer = tokenizer.decode(
+        outputs[0][input_length:],
         skip_special_tokens=True,
-
     )
 
 
-
-    return text.strip()
+    return answer.strip()
 
 
 
 # ============================================================
-# HEALTH
-# ============================================================
-
-
-def health():
-
-    return {
-
-        "status": "ok",
-
-        "model": MODEL_ID,
-
-        "device": DEVICE,
-
-        "cuda": torch.cuda.is_available(),
-
-        "timestamp": time.time(),
-
-    }
-
-
-
-# ============================================================
-# GRADIO UI
+# UI
 # ============================================================
 
 
 with gr.Blocks(
-    title="X-RUDRA Model A"
+    title="X-RUDRA M2"
 ) as demo:
 
 
     gr.Markdown(
         f"""
-# ⚡ X-RUDRA — Model A
-
-### WeiboAI/VibeThinker-3B
+# ⚡ X-RUDRA M2
 
-Independent reasoning model.
-
-**Model ID**
+Model:
 
 `{MODEL_ID}`
 
-**Device**
+Device:
 
 `{DEVICE}`
 """
     )
 
 
-
     prompt = gr.Textbox(
-
         label="Prompt",
-
-        placeholder=(
-            "Ask anything..."
-        ),
-
-        lines=8,
-
+        lines=6,
+        placeholder="Ask something..."
     )
 
 
-
     with gr.Row():
 
-
         max_tokens = gr.Slider(
-
-            minimum=128,
-
-            maximum=4096,
-
-            value=1024,
-
-            step=128,
-
-            label="Max tokens",
-
+            64,
+            2048,
+            value=512,
+            step=64,
+            label="Max Tokens"
         )
 
 
         temperature = gr.Slider(
-
-            minimum=0,
-
-            maximum=1.5,
-
+            0.1,
+            1.5,
             value=0.7,
-
-            step=0.05,
-
-            label="Temperature",
-
+            step=0.1,
+            label="Temperature"
         )
 
 
-        top_p = gr.Slider(
-
-            minimum=0.1,
-
-            maximum=1.0,
-
-            value=0.9,
-
-            step=0.05,
-
-            label="Top P",
-
-        )
-
-
-
-    button = gr.Button(
-
-        "🚀 Generate",
-
-        variant="primary",
-
+    btn = gr.Button(
+        "Generate",
+        variant="primary"
     )
 
 
-
-    output = gr.Textbox(
-
-        label="Response",
-
-        lines=20,
-
+    output = gr.Markdown(
+        label="Response"
     )
 
 
-
-    button.click(
-
+    btn.click(
         generate,
-
         inputs=[
-
             prompt,
-
             max_tokens,
-
             temperature,
-
-            top_p,
-
         ],
-
         outputs=output,
-
-        api_name="generate",
-
-    )
-
-
-
-    gr.JSON(
-
-        value=health(),
-
-        label="Health",
-
     )
 
 
-
 # ============================================================
 # START
 # ============================================================
 
-
 if __name__ == "__main__":
 
 
     demo.launch(
-
         server_name="0.0.0.0",
-
-        server_port=int(
-
-            os.getenv(
-                "PORT",
-                "7860"
-            )
-
-        ),
-
+        server_port=7860,
     )