Update app.py

app.py

@@ -19,7 +19,7 @@ pipeline = transformers.pipeline(
 tokenizer = AutoTokenizer.from_pretrained(model)
 
 
-def predict_answer(question):
+def predict_answer(question, token=25):
     
     messages = [{"role": "user", "content": f"{question}"}]
     
@@ -27,20 +27,21 @@ def predict_answer(question):
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     
     
-    outputs = pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+    outputs = pipeline(prompt, max_new_tokens=token, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
     
     return outputs[0]["generated_text"]
 
 
 
-def gradio_predict(question):
-    answer = predict_answer(question)
+def gradio_predict(question, token):
+    answer = predict_answer(question, token)
     return answer
 
 # Define the Gradio interface
 iface = gr.Interface(
     fn=gradio_predict,
-    inputs=[gr.Textbox(label="Question", placeholder="e.g. What are the colors of the bus in the image?", scale=4)],
+    inputs=[gr.Textbox(label="Question", placeholder="e.g. What are the colors of the bus in the image?", scale=4),
+           gr.Slider(2, 100, value=25, label="Token Count", info="Choose between 2 and 100")]],
     outputs=gr.TextArea(label="Answer"),
     title="KatzBot",
     description="Llama - Mistral Merge",