Update app.py

app.py

@@ -1,8 +1,7 @@
 import os
-import time
 import torch
-import gradio as gr
 import spaces
+import gradio as gr
 
 from transformers import (
     AutoTokenizer,
@@ -16,70 +15,54 @@ from transformers import (
 
 MODEL_ID = os.getenv(
     "MODEL_ID",
-    "LiquidAI/LFM2-2.6B",
+    "LiquidAI/LFM2-2.6B"
+)
+
+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)
 
 
 # ============================================================
-# LAZY MODEL LOADING
+# 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,
-    )
-
-    kwargs = {
-        "trust_remote_code": True,
-    }
-
-
-    if torch.cuda.is_available():
-
-        kwargs["device_map"] = "auto"
-        kwargs["torch_dtype"] = torch.bfloat16
-
-    else:
-
-        kwargs["torch_dtype"] = torch.float32
+print("Loading tokenizer...")
 
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+)
 
-    print("Loading model...")
 
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        **kwargs,
-    )
+print("Loading model...")
 
-    model.eval()
 
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=(
+        torch.float16
+        if DEVICE == "cuda"
+        else torch.float32
+    ),
+    device_map="auto",
+    trust_remote_code=True,
+)
 
-    if tokenizer.pad_token_id is None:
 
-        tokenizer.pad_token = tokenizer.eos_token
+model.eval()
 
 
-    print("MODEL READY")
+print("MODEL READY")
 
 
 # ============================================================
@@ -87,268 +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 or not prompt.strip():
-
-        return "ERROR: Empty prompt"
-
-
-    load_model()
+    if not prompt.strip():
 
+        return "Enter prompt."
 
-    messages = [
-        {
-            "role": "user",
-            "content": prompt.strip(),
-        }
-    ]
 
-
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
+    inputs = tokenizer(
+        prompt,
         return_tensors="pt",
+        padding=True,
+        truncation=True,
     )
 
 
-    if torch.cuda.is_available():
-
-        inputs = inputs.to(
-            model.device
-        )
-
-    else:
-
-        inputs = inputs.to(
-            DEVICE
-        )
+    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()
-
-
-
-# ============================================================
-# HEALTH
-# ============================================================
-
-
-def health():
-
-    return {
-        "status": "ok",
-        "model": MODEL_ID,
-        "device": DEVICE,
-        "cuda": torch.cuda.is_available(),
-        "timestamp": time.time(),
-    }
+    return answer.strip()
 
 
 
 # ============================================================
-# GRADIO UI
+# UI
 # ============================================================
 
 
 with gr.Blocks(
-    title="X-RUDRA M1/M2"
+    title="X-RUDRA M2"
 ) as demo:
 
 
     gr.Markdown(
-        """
-# ⚡ X-RUDRA Model Space
-
-Independent reasoning model endpoint.
+        f"""
+# ⚡ X-RUDRA M2
 
 Model:
-`{}`
+
+`{MODEL_ID}`
 
 Device:
-`{}`
-""".format(
-            MODEL_ID,
-            DEVICE,
-        )
+
+`{DEVICE}`
+"""
     )
 
 
     prompt = gr.Textbox(
-
         label="Prompt",
-
-        placeholder=(
-            "Enter your question..."
-        ),
-
-        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",
-        )
-
-
-        top_p = gr.Slider(
-
-            minimum=0.1,
-
-            maximum=1,
-
-            value=0.9,
-
-            step=0.05,
-
-            label="Top P",
+            step=0.1,
+            label="Temperature"
         )
 
 
-
-    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(
-
-        fn=generate,
-
+    btn.click(
+        generate,
         inputs=[
-
             prompt,
-
             max_tokens,
-
             temperature,
-
-            top_p,
-
         ],
-
         outputs=output,
-
-        api_name="generate",
-    )
-
-
-    gr.JSON(
-        health(),
-        label="Health"
     )
 
 
-
 # ============================================================
 # START
 # ============================================================
 
-
 if __name__ == "__main__":
 
-    demo.launch(
 
+    demo.launch(
         server_name="0.0.0.0",
-
-        server_port=int(
-
-            os.getenv(
-                "PORT",
-                7860
-            )
-
-        ),
+        server_port=7860,
     )