Update app.py

app.py

@@ -24,7 +24,7 @@ model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloa
 @torch.inference_mode()
 @spaces.GPU
 def predict_math_bot(user_message, system_message="", max_new_tokens=125, temperature=0.1, top_p=0.9, repetition_penalty=1.9, do_sample=False):
-    prompt = f"<|system|>\n{system_message}\n<|user|>\n{user_message}<|assistant|>" if system_message else user_message
+    prompt = f"<s><INST>{user_message}{system_message}<|\INST|>" if system_message else user_message
     inputs = tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
     input_ids = inputs["input_ids"].to(model.device)
 
@@ -53,16 +53,16 @@ def main():
                 top_p = gr.Slider(label="Top-p (nucleus sampling)", value=0.90, minimum=0.01, maximum=0.99)
                 repetition_penalty = gr.Slider(label="Repetition penalty", value=1.9, minimum=1.0, maximum=2.0)
                 do_sample = gr.Checkbox(label="Uncheck for faster inference", value=False)
-        
-        with gr.Row():
-            user_message = gr.Textbox(label="🫡Your Message", lines=3, placeholder="Enter your math query here...")
-            system_message = gr.Textbox(label="📉System Prompt", lines=2, placeholder="Optional: Set a scene or introduce a character...")
-
         gr.Button("Try🫡📉MetaMath").click(
             predict_math_bot,
             inputs=[user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample],
             outputs=output_text
-        )
+        )     
+        with gr.Row():
+            user_message = gr.Textbox(label="🫡Your Message", lines=3, placeholder="Enter your math query here...")
+            system_message = gr.Textbox(label="📉System Prompt", lines=2, placeholder="Optional: Set a scene or introduce a character...")
+
+
 
         # Example placeholders
 #       demo.add_examples([