Update app.py

app.py

@@ -29,7 +29,7 @@ class TuluChatBot:
         prompt = f"<|assistant|>\n {self.system_message}\n\n <|user|>{user_message}\n\n<|assistant|>\n"
         return prompt
 
-    def predict(self, user_message, temperature=0.4, max_new_tokens=70, top_p=0.99, repetition_penalty=1.9):
+    def predict(self, user_message, temperature, max_new_tokens, top_p, repetition_penalty, do_sample):
         prompt = self.format_prompt(user_message)
         inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
         input_ids = inputs["input_ids"].to(self.model.device)
@@ -43,16 +43,21 @@ class TuluChatBot:
             temperature=temperature,
             top_p=top_p,
             repetition_penalty=repetition_penalty,
-            do_sample=True
+            do_sample=do_sample
         )
 
         response = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
         return response
 
-def gradio_predict(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty):
+def gradio_predict(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty, do_sample):
     Tulu_bot.set_system_message(system_message)
-    response = Tulu_bot.predict(user_message, temperature, max_new_tokens, top_p, repetition_penalty)
+    if not do_sample:
+        temperature = 1.2  # Default value
+        top_p = 0.9        # Default value
+        repetition_penalty = 0.9  # Default value
+    response = Tulu_bot.predict(user_message, temperature, max_new_tokens, top_p, repetition_penalty, do_sample)
     return response
+
     
 Tulu_bot = TuluChatBot(model, tokenizer)
 
@@ -63,10 +68,11 @@ iface = gr.Interface(
     inputs=[
         gr.Textbox(label="Your Message", type="text", lines=3),
         gr.Textbox(label="Introduce a Character Here or Set a Scene (system prompt)", type="text", lines=2),
-        gr.Slider(label="Max new tokens", value=1269, minimum=550, maximum=3200, step=1),
-        gr.Slider(label="Temperature", value=1.2, minimum=0.05, maximum=4.0, step=0.05),
-        gr.Slider(label="Top-p (nucleus sampling)", value=0.90, minimum=0.01, maximum=0.99, step=0.05),
-        gr.Slider(label="Repetition penalty", value=1.9, minimum=1.0, maximum=2.0, step=0.05)
+        gr.Checkbox(label="Advanced", value=True, id="do_sample"),  # Add an id to the checkbox
+        gr.Slider(label="Max new tokens", value=1269, minimum=550, maximum=3200, step=1, visible=gr.Visibility(id="do_sample", value=True)),
+        gr.Slider(label="Temperature", value=1.2, minimum=0.05, maximum=4.0, step=0.05, visible=gr.Visibility(id="do_sample", value=True)),
+        gr.Slider(label="Top-p (nucleus sampling)", value=0.90, minimum=0.01, maximum=0.99, step=0.05, visible=gr.Visibility(id="do_sample", value=True)),
+        gr.Slider(label="Repetition penalty", value=1.9, minimum=1.0, maximum=2.0, step=0.05, visible=gr.Visibility(id="do_sample", value=True))
     ],
     outputs="text",
     theme="ParityError/Anime"