Update app.py

app.py

@@ -11,10 +11,11 @@ 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
+# System prompt – use triple quotes for multi-line strings
+SYSTEM_PROMPT = """
+You are X-RUDRA, a helpful, knowledgeable, and concise AI assistant.
+
+# SEARCH DECISION SYSTEM PROMPT
 
 ACT ONLY WHEN REQUIRED.
 
@@ -106,14 +107,13 @@ 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 + API)")
+print("X-RUDRA M1 (CHAT + API)")   # Change to M2 for the other Space
 print("MODEL:", MODEL_ID)
 print("DEVICE:", DEVICE)
-print("SYSTEM PROMPT:", SYSTEM_PROMPT)
+print("SYSTEM PROMPT (first 100 chars):", SYSTEM_PROMPT[:100] + "...")
 print("=" * 60)
 
 # ============================================================
@@ -148,14 +148,12 @@ def build_prompt_with_system(history, new_user_message=None):
     new_user_message: str (if provided, appended as user message)
     Returns: prompt string ready for tokenization.
     """
-    # 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})
 
     # 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(
@@ -174,7 +172,6 @@ def build_prompt_with_system(history, new_user_message=None):
             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
 
@@ -185,16 +182,11 @@ def build_prompt_with_system(history, new_user_message=None):
 
 @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(
         prompt,
         return_tensors="pt",
@@ -222,7 +214,6 @@ 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 user message and assistant response to history
     history.append({"role": "user", "content": message})
     history.append({"role": "assistant", "content": answer})
 
@@ -235,12 +226,6 @@ def generate_response(message, history, max_tokens, temperature):
 
 @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)
 
@@ -277,13 +262,12 @@ def generate(prompt, max_tokens, temperature):
 # UI – Chat Interface
 # ============================================================
 
-with gr.Blocks(title="X-RUDRA M2") as demo:
+with gr.Blocks(title="X-RUDRA M1") as demo:   # Change to M2 for M2 Space
     gr.Markdown(
         f"""
-# ⚡ X-RUDRA M2 – Chat + API
+# ⚡ X-RUDRA M1 – Chat + API
 **Model:** `{MODEL_ID}`  
 **Device:** `{DEVICE}`  
-**System Prompt:** `{SYSTEM_PROMPT[:80]}...`
 """
     )
 
@@ -307,9 +291,7 @@ with gr.Blocks(title="X-RUDRA M2") as demo:
         outputs=[msg, chatbot]
     )
 
-    # ------------------------------------------------------------
-    # Hidden Interface for API – exposes /generate endpoint
-    # ------------------------------------------------------------
+    # Hidden API endpoint
     gr.Interface(
         fn=generate,
         inputs=[
@@ -318,10 +300,10 @@ with gr.Blocks(title="X-RUDRA M2") as demo:
             gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="temperature")
         ],
         outputs=gr.Textbox(label="response"),
-        title="X-RUDRA M2 API",
+        title="X-RUDRA M1 API",
         description="Standalone generation endpoint.",
         api_name="generate",
-        visible=False,  # Hidden from UI, but API route is still active
+        visible=False,
     )