Update app.py

app.py

@@ -36,8 +36,11 @@ model.eval()
 def generate(
     message: str,
     chat_history: list[dict],
-    max_new_tokens: int = 1048,
-    temperature: float = 0.01,
+    max_new_tokens: int = 1024,
+    temperature: float = 0.6,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
     conversation = chat_history.copy()
     conversation.append({"role": "user", "content": message})
@@ -54,8 +57,11 @@ def generate(
         streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
         temperature=temperature,
         num_beams=1,
+        repetition_penalty=repetition_penalty
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
@@ -69,12 +75,40 @@ def generate(
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
+        gr.Slider(
+            label="Max new tokens",
+            minimum=1,
+            maximum=MAX_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        ),
         gr.Slider(
             label="Temperature",
             minimum=0.1,
             maximum=4.0,
             step=0.1,
-            value=0.01,
+            value=0.6,
+        ),
+        gr.Slider(
+            label="Top-p (nucleus sampling)",
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.9,
+        ),
+        gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=1000,
+            step=1,
+            value=50,
+        ),
+        gr.Slider(
+            label="Repetition penalty",
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            value=1.2,
         ),
     ],
     stop_btn=None,