Update app.py

app.py

@@ -13,13 +13,13 @@ llm2 = Llama(model_path=filename, seed=random.randint(1, 2**31))
 title = """<h1 align="center">Chat with awesome WizardLM 7b model!</h1><br>"""
 description = "This model is awesome for its size! It is only 20th the size of Chatgpt but is around 90% as good as Chatgpt. However, please don't rely on WizardLM to provide 100% true information as it might be wrong sometimes."
 
-def bot(user_message, temperature, top_p, repeat_penalty):
+def bot(user_message, temperature, top_p, top_k, repeat_penalty):
     tokens3 = llm2.tokenize(user_message.encode())
     token4 = llm2.tokenize(b"\n\n### Response:")
     tokens = tokens3 + token4
     
     output = ""
-    for token in llm2.generate(tokens, top_k=50, top_p=top_p, temp=temperature, repeat_penalty=repeat_penalty):
+    for token in llm2.generate(tokens, top_k=top_k, top_p=top_p, temp=temperature, repeat_penalty=repeat_penalty):
         text = llm2.detokenize([token])
         output += text.decode()
         if token == llm2.token_eos():
@@ -35,6 +35,7 @@ interface = gr.Interface(
         gr.Textbox(label="Your Message", placeholder="Type your message here..."),
         gr.Slider(minimum=0, maximum=2, default=1, label="Temperature"),
         gr.Slider(minimum=0, maximum=1, default=0.73, label="Top P"),
+        gr.Slider(minimum=0, maximum=2048, default=50, label="Top_K"),
         gr.Slider(minimum=0, maximum=2, default=1.1, label="Repeat Penalty")
     ],
     outputs="text",