Update app.py

app.py

@@ -11,10 +11,109 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 MODEL_ID = os.getenv("MODEL_ID", "WeiboAI/VibeThinker-3B")
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+# System prompt – defines the assistant's personality and constraints
+SYSTEM_PROMPT = (
+    "You are X-RUDRA, a helpful, knowledgeable, and concise AI assistant. "
+    "# SEARCH DECISION SYSTEM PROMPT
+
+ACT ONLY WHEN REQUIRED.
+
+SEARCH WHEN:
+
+* SEARCH
+* BROWSE
+* LOOKUP
+* VERIFY
+* CHECK
+* FIND
+* RESEARCH
+* COMPARE CURRENT DATA
+* CONFIRM LATEST DATA
+* RETRIEVE EXTERNAL INFORMATION
+* HANDLE UNCERTAIN FACTS
+* HANDLE TIME-SENSITIVE INFORMATION
+* HANDLE NICHE INFORMATION
+* HANDLE LOCAL INFORMATION
+* HANDLE CURRENT PRICES
+* HANDLE CURRENT NEWS
+* HANDLE CURRENT SPORTS
+* HANDLE CURRENT PRODUCTS
+* HANDLE CURRENT PEOPLE
+* HANDLE CURRENT COMPANIES
+* HANDLE CURRENT SOFTWARE
+* HANDLE CURRENT DOCUMENTATION
+
+DO NOT SEARCH WHEN:
+
+* CHAT
+* CONVERSE
+* GREET
+* JOKE
+* BRAINSTORM
+* EXPLAIN FROM KNOWN KNOWLEDGE
+* REWRITE
+* TRANSLATE
+* SUMMARIZE PROVIDED TEXT
+* WRITE
+* CODE FROM PROVIDED REQUIREMENTS
+* SOLVE SIMPLE REASONING
+* ANSWER CASUAL QUESTIONS
+* HANDLE TIMEPASS CONVERSATION
+
+PRIORITIZE:
+
+* USER INTENT
+* ACCURACY
+* FRESHNESS
+* RELEVANCE
+* PRIMARY SOURCES
+* OFFICIAL SOURCES
+* DIRECT EVIDENCE
+
+AVOID:
+
+* UNNECESSARY SEARCHES
+* SEARCHING CASUAL CONVERSATION
+* SEARCHING EVERY MESSAGE
+* FABRICATING SEARCH RESULTS
+* FABRICATING SOURCES
+* FABRICATING CITATIONS
+* USING OUTDATED INFORMATION WHEN FRESH INFORMATION IS REQUIRED
+
+WHEN SEARCHING:
+
+1. IDENTIFY THE INFORMATION REQUIRED.
+2. FORMULATE PRECISE QUERIES.
+3. SEARCH RELEVANT SOURCES.
+4. VERIFY IMPORTANT CLAIMS.
+5. PREFER PRIMARY SOURCES.
+6. CROSS-CHECK CONFLICTING INFORMATION.
+7. DISTINGUISH FACT FROM INFERENCE.
+8. CITE SOURCES.
+9. ANSWER DIRECTLY.
+10. STOP SEARCHING WHEN SUFFICIENT EVIDENCE EXISTS.
+
+WHEN NOT SEARCHING:
+
+1. UNDERSTAND THE REQUEST.
+2. USE AVAILABLE CONTEXT.
+3. ANSWER DIRECTLY.
+4. DO NOT PERFORM A SEARCH JUST TO APPEAR HELPFUL.
+
+CORE RULE:
+
+SEARCH FOR INFORMATION.
+DO NOT SEARCH FOR CONVERSATION.
+
+SEARCH ONLY WHEN SEARCHING IMPROVES ACCURACY, FRESHNESS, VERIFICATION, OR COMPLETENESS.
+"
+)
+
 print("=" * 60)
-print("X-RUDRA M2 (CHAT)")
+print("X-RUDRA M1 (CHAT + API)")
 print("MODEL:", MODEL_ID)
 print("DEVICE:", DEVICE)
+print("SYSTEM PROMPT:", SYSTEM_PROMPT)
 print("=" * 60)
 
 # ============================================================
@@ -30,52 +129,70 @@ if tokenizer.pad_token is None:
 print("Loading model...")
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    dtype=torch.float16 if DEVICE == "cuda" else torch.float32,   # FIXED
+    dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
     device_map="auto",
     trust_remote_code=True,
 )
 model.eval()
 print("MODEL READY")
 
+
 # ============================================================
-# GENERATION FUNCTION (history = list of dicts)
+# HELPER: BUILD PROMPT WITH SYSTEM + HISTORY
 # ============================================================
 
-@spaces.GPU
-def generate_response(message, history, max_tokens, temperature):
+def build_prompt_with_system(history, new_user_message=None):
     """
-    history: list of dicts with 'role' and 'content' (e.g. [{"role":"user","content":"Hi"}])
-    returns: (new_message, updated_history) – new_message = '' to clear input.
+    Build a full prompt string from conversation history and an optional new user message.
+    history: list of dicts with 'role' and 'content' (user/assistant)
+    new_user_message: str (if provided, appended as user message)
+    Returns: prompt string ready for tokenization.
     """
-    # Ensure history is always a list
-    if history is None:
-        history = []
-
-    # Append the new user message
-    history.append({"role": "user", "content": message})
+    # Create a copy of history and optionally add the new user message
+    messages = list(history) if history else []
+    if new_user_message is not None:
+        messages.append({"role": "user", "content": new_user_message})
 
-    # Build prompt using chat template (if available)
-    prompt = None
+    # If the tokenizer has a chat template that supports system, use it
     if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
+        # Some templates expect a system message; we'll include it
+        full_messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
         try:
             prompt = tokenizer.apply_chat_template(
-                history,
+                full_messages,
                 tokenize=False,
                 add_generation_prompt=True
             )
+            return prompt
         except Exception as e:
             print("Chat template failed, falling back to manual format:", e)
 
-    # Fallback manual format
-    if prompt is None:
-        manual = ""
-        for turn in history:
-            if turn["role"] == "user":
-                manual += f"User: {turn['content']}\n"
-            elif turn["role"] == "assistant":
-                manual += f"Assistant: {turn['content']}\n"
-        manual += "Assistant:"
-        prompt = manual
+    # Fallback: manual formatting with system prompt
+    prompt = f"System: {SYSTEM_PROMPT}\n"
+    for turn in messages:
+        if turn["role"] == "user":
+            prompt += f"User: {turn['content']}\n"
+        elif turn["role"] == "assistant":
+            prompt += f"Assistant: {turn['content']}\n"
+    # Add a final "Assistant:" to prompt the model
+    prompt += "Assistant:"
+    return prompt
+
+
+# ============================================================
+# GENERATION FUNCTION (for chat UI)
+# ============================================================
+
+@spaces.GPU
+def generate_response(message, history, max_tokens, temperature):
+    """
+    Takes the current message and history, returns updated history with assistant reply.
+    """
+    if history is None:
+        history = []
+
+    # Build prompt including the new user message
+    prompt = build_prompt_with_system(history, message)
 
     # Tokenize
     inputs = tokenizer(
@@ -86,7 +203,6 @@ def generate_response(message, history, max_tokens, temperature):
         padding=True,
     )
     inputs = {k: v.to(model.device) for k, v in inputs.items()}
-
     input_len = inputs["input_ids"].shape[-1]
 
     with torch.no_grad():
@@ -106,29 +222,72 @@ def generate_response(message, history, max_tokens, temperature):
     new_tokens = outputs[0][input_len:]
     answer = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
-    # Append assistant response to history
+    # Append user message and assistant response to history
+    history.append({"role": "user", "content": message})
     history.append({"role": "assistant", "content": answer})
 
-    # Return empty string to clear the input box and the updated history
     return "", history
 
+
 # ============================================================
-# UI – Chatbot WITHOUT 'type' argument (defaults to "messages")
+# GENERATION FUNCTION (for API – standalone)
 # ============================================================
 
-with gr.Blocks(title="X-RUDRA M2") as demo:
+@spaces.GPU
+def generate(prompt, max_tokens, temperature):
+    """
+    Standalone generation for API calls.
+    Expects a raw prompt string, returns the generated text.
+    """
+    # Build full prompt with system + user input
+    # We treat the input as a user message
+    messages = [{"role": "user", "content": prompt}]
+    full_prompt = build_prompt_with_system(messages)
+
+    inputs = tokenizer(
+        full_prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=4096,
+        padding=True,
+    )
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    input_len = inputs["input_ids"].shape[-1]
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=int(max_tokens),
+            temperature=float(temperature),
+            do_sample=True,
+            top_p=0.95,
+            top_k=50,
+            repetition_penalty=1.15,
+            no_repeat_ngram_size=3,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    new_tokens = outputs[0][input_len:]
+    answer = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+    return answer
+
+
+# ============================================================
+# UI – Chat Interface
+# ============================================================
+
+with gr.Blocks(title="X-RUDRA M1") as demo:
     gr.Markdown(
         f"""
-# ⚡ X-RUDRA M2 – Chat
+# ⚡ X-RUDRA M1 – Chat + API
 **Model:** `{MODEL_ID}`  
-**Device:** `{DEVICE}`
+**Device:** `{DEVICE}`  
+**System Prompt:** `{SYSTEM_PROMPT[:80]}...`
 """
     )
 
-    # DO NOT add 'type=' here – let Gradio use its default.
-    # The default on this version is "messages" (list of dicts).
     chatbot = gr.Chatbot(height=600, label="Conversation")
-
     with gr.Row():
         msg = gr.Textbox(placeholder="Ask anything...", scale=8)
         send = gr.Button("Send", variant="primary", scale=1)
@@ -148,6 +307,24 @@ with gr.Blocks(title="X-RUDRA M2") as demo:
         outputs=[msg, chatbot]
     )
 
+    # ------------------------------------------------------------
+    # Hidden Interface for API – exposes /generate endpoint
+    # ------------------------------------------------------------
+    gr.Interface(
+        fn=generate,
+        inputs=[
+            gr.Textbox(label="prompt", lines=2),
+            gr.Slider(64, 2048, value=512, step=64, label="max_tokens"),
+            gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="temperature")
+        ],
+        outputs=gr.Textbox(label="response"),
+        title="X-RUDRA M1 API",
+        description="Standalone generation endpoint.",
+        api_name="generate",
+        visible=False,  # Hidden from UI, but API route is still active
+    )
+
+
 # ============================================================
 # START
 # ============================================================