Update app.py

app.py

@@ -2,9 +2,17 @@ import os
 import time
 import torch
 import gradio as gr
+import spaces
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
 
-from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# ============================================================
+# CONFIG
+# ============================================================
 
 MODEL_ID = os.getenv(
     "MODEL_ID",
@@ -13,45 +21,89 @@ MODEL_ID = os.getenv(
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-print(f"Loading model: {MODEL_ID}")
-print(f"Device: {DEVICE}")
+print("X-RUDRA MODEL SPACE")
+print("MODEL:", MODEL_ID)
+print("DEVICE:", DEVICE)
 
 
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-)
+# ============================================================
+# LAZY MODEL LOADING
+# ============================================================
 
-model_kwargs = {
-    "trust_remote_code": True,
-}
+tokenizer = None
+model = None
 
-if torch.cuda.is_available():
-    model_kwargs["device_map"] = "auto"
-    model_kwargs["torch_dtype"] = torch.bfloat16
-else:
-    model_kwargs["torch_dtype"] = torch.float32
 
+def load_model():
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    **model_kwargs,
-)
+    global tokenizer
+    global model
+
+    if model is not None:
+        return
+
+    print("Loading tokenizer...")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True,
+    )
 
-model.eval()
+    kwargs = {
+        "trust_remote_code": True,
+    }
+
+
+    if torch.cuda.is_available():
+
+        kwargs["device_map"] = "auto"
+        kwargs["torch_dtype"] = torch.bfloat16
+
+    else:
 
-if tokenizer.pad_token_id is None:
-    tokenizer.pad_token = tokenizer.eos_token
+        kwargs["torch_dtype"] = torch.float32
 
 
+    print("Loading model...")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        **kwargs,
+    )
+
+    model.eval()
+
+
+    if tokenizer.pad_token_id is None:
+
+        tokenizer.pad_token = tokenizer.eos_token
+
+
+    print("MODEL READY")
+
+
+# ============================================================
+# GENERATION
+# ============================================================
+
+
+@spaces.GPU(
+    duration=120
+)
 def generate(
-    prompt: str,
-    max_new_tokens: int = 1024,
-    temperature: float = 0.7,
-    top_p: float = 0.9,
+    prompt,
+    max_new_tokens,
+    temperature,
+    top_p,
 ):
+
     if not prompt or not prompt.strip():
-        return "ERROR: prompt is empty."
+
+        return "ERROR: Empty prompt"
+
+
+    load_model()
+
 
     messages = [
         {
@@ -60,6 +112,7 @@ def generate(
         }
     ]
 
+
     inputs = tokenizer.apply_chat_template(
         messages,
         tokenize=True,
@@ -67,35 +120,77 @@ def generate(
         return_tensors="pt",
     )
 
+
     if torch.cuda.is_available():
-        inputs = inputs.to(model.device)
+
+        inputs = inputs.to(
+            model.device
+        )
+
     else:
-        inputs = inputs.to(DEVICE)
+
+        inputs = inputs.to(
+            DEVICE
+        )
+
 
     input_length = inputs.shape[-1]
 
+
     with torch.inference_mode():
+
         output = model.generate(
+
             inputs,
-            max_new_tokens=int(max_new_tokens),
-            temperature=float(temperature),
-            top_p=float(top_p),
-            do_sample=float(temperature) > 0,
-            pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id,
+
+            max_new_tokens=int(
+                max_new_tokens
+            ),
+
+            temperature=float(
+                temperature
+            ),
+
+            top_p=float(
+                top_p
+            ),
+
+            do_sample=(
+                float(temperature) > 0
+            ),
+
+            pad_token_id=(
+                tokenizer.pad_token_id
+            ),
+
+            eos_token_id=(
+                tokenizer.eos_token_id
+            ),
         )
 
-    generated = output[0][input_length:]
+
+    generated = output[0][
+        input_length:
+    ]
+
 
     text = tokenizer.decode(
         generated,
         skip_special_tokens=True,
     )
 
+
     return text.strip()
 
 
+
+# ============================================================
+# HEALTH
+# ============================================================
+
+
 def health():
+
     return {
         "status": "ok",
         "model": MODEL_ID,
@@ -105,76 +200,155 @@ def health():
     }
 
 
-with gr.Blocks(title="X RUDRA - LFM2 2.6B") as demo:
+
+# ============================================================
+# GRADIO UI
+# ============================================================
+
+
+with gr.Blocks(
+    title="X-RUDRA M1/M2"
+) as demo:
+
 
     gr.Markdown(
         """
-        # ⚡ X RUDRA — Model B
+# ⚡ X-RUDRA Model Space
 
-        **LiquidAI/LFM2-2.6B**
+Independent reasoning model endpoint.
 
-        Independent second-model reasoning endpoint.
-        """
+Model:
+`{}`
+
+Device:
+`{}`
+""".format(
+            MODEL_ID,
+            DEVICE,
+        )
     )
 
+
     prompt = gr.Textbox(
+
         label="Prompt",
-        placeholder="Enter a question...",
+
+        placeholder=(
+            "Enter your question..."
+        ),
+
         lines=8,
     )
 
+
     with gr.Row():
+
+
         max_tokens = gr.Slider(
-            128,
-            4096,
+
+            minimum=128,
+
+            maximum=4096,
+
             value=1024,
+
             step=128,
-            label="Max new tokens",
+
+            label="Max Tokens",
         )
 
+
         temperature = gr.Slider(
-            0.0,
-            1.5,
+
+            minimum=0,
+
+            maximum=1.5,
+
             value=0.7,
+
             step=0.05,
+
             label="Temperature",
         )
 
+
         top_p = gr.Slider(
-            0.1,
-            1.0,
+
+            minimum=0.1,
+
+            maximum=1,
+
             value=0.9,
+
             step=0.05,
+
             label="Top P",
         )
 
-    generate_button = gr.Button(
-        "Generate",
+
+
+    button = gr.Button(
+
+        "🚀 Generate",
+
         variant="primary",
     )
 
+
     output = gr.Textbox(
+
         label="Response",
+
         lines=20,
     )
 
-    generate_button.click(
-        generate,
+
+    button.click(
+
+        fn=generate,
+
         inputs=[
+
             prompt,
+
             max_tokens,
+
             temperature,
+
             top_p,
+
         ],
+
         outputs=output,
+
         api_name="generate",
     )
 
-    gr.api(
-        health,
-        api_name="health",
+
+    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
+            )
+
+        ),
+    )