added top_k

app.py

@@ -51,6 +51,7 @@ def respond(
     max_tokens,
     temperature,
     top_p,
+    top_k,
     repetition_penalty
 ):
     messages = [{"role": "system", "content": system_message}]
@@ -66,7 +67,7 @@ def respond(
     streamer = CustomIterable()
 
     inputs = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt", add_generation_prompt=True)
-    thread = threading.Thread(target=model.generate, args=([inputs]), kwargs={"max_new_tokens": max_tokens, "do_sample": True, "temperature": temperature, "top_p": top_p, "repetition_penalty": repetition_penalty, "streamer": streamer})
+    thread = threading.Thread(target=model.generate, args=([inputs]), kwargs={"max_new_tokens": max_tokens, "do_sample": True, "temperature": temperature, "top_p": top_p, "top_k": top_k, "repetition_penalty": repetition_penalty, "streamer": streamer})
     thread.start()
     response = ""
 
@@ -84,7 +85,7 @@ demo = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Textbox(value="Du bist ein hilfreicher Assistent.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=1, maximum=1024, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.3, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
@@ -93,6 +94,13 @@ demo = gr.ChatInterface(
             step=0.05,
             label="Top-p (nucleus sampling)",
         ),
+        gr.Slider(
+            minimum=16,
+            maximum=1024,
+            value=512,
+            step=1,
+            label="Top-k",
+        ),
         gr.Slider(
             minimum=0.1,
             maximum=2.0,