Update app.py

app.py

@@ -18,13 +18,13 @@ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder
 model_name = 'TencentARC/Mistral_Pro_8B_v0.1'
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
-model.generation_config = GenerationConfig.from_pretrained(model_name)
+# model.generation_config = GenerationConfig.from_pretrained(model_name)
 model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
 @torch.inference_mode()
 @spaces.GPU
-def predict_math_bot(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty):
-    prompt = f"<|system|>\n{self.system_message}\n<|user|>\n{user_message}<|assistant|>" if system_message else user_message
+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
     inputs = tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
     input_ids = inputs["input_ids"].to(model.device)
 
@@ -35,7 +35,7 @@ def predict_math_bot(user_message, system_message, max_new_tokens, temperature,
         top_p=top_p,
         repetition_penalty=repetition_penalty,
         pad_token_id=tokenizer.eos_token_id,
-        do_sample=True
+        do_sample=do_sample
     )
 
     response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
@@ -52,15 +52,15 @@ def main():
                 temperature = gr.Slider(label="Temperature", value=0.1, minimum=0.05, maximum=1.0)
                 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)
-                use_custom_settings = gr.Checkbox(label="Use custom settings", value=False)
+                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("Generate").click(
+        gr.Button("Try🫡📉MetaMath").click(
             predict_math_bot,
-            inputs=[user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, use_custom_settings],
+            inputs=[user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample],
             outputs=output_text
         )