Included params on ui

app.py

@@ -1,12 +1,15 @@
 import os
 import gradio as gr
-from gradio.components import Textbox, Button
+from gradio.components import Textbox, Button, Slider, Checkbox
 from AinaTheme import theme
 from urllib.error import HTTPError
 
 from rag import RAG
 from utils import setup
 
+MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", default=100))
+SHOW_MODEL_PARAMETERS_IN_UI = os.environ.get("SHOW_MODEL_PARAMETERS_IN_UI", default="True") == "True"
+
 setup()
 
 
@@ -19,9 +22,9 @@ rag = RAG(
 )
 
 
-def generate(prompt):
+def generate(prompt, model_parameters):
     try:
-        output = rag.get_response(prompt)
+        output = rag.get_response(prompt, model_parameters)
         return output
     except HTTPError as err:
         if err.code == 400:
@@ -34,12 +37,23 @@ def generate(prompt):
         )
 
 
-def submit_input(input_):
+def submit_input(input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature):
     if input_.strip() == "":
         gr.Warning("Not possible to inference an empty input")
         return None
 
-    output = generate(input_)
+
+    model_parameters = {
+        "max_new_tokens": max_new_tokens,
+        "repetition_penalty": repetition_penalty,
+        "top_k": top_k,
+        "top_p": top_p,
+        "do_sample": do_sample,
+        "num_beams": num_beams,
+        "temperature": temperature
+    }
+
+    output = generate(input_, model_parameters)
 
     return output
 
@@ -52,8 +66,15 @@ def change_interactive(text):
 
 def clear():
     return (
+        None, 
         None,
-        None,
+        gr.Slider(value=100),
+        gr.Slider(value=1.2),
+        gr.Slider(value=50),
+        gr.Slider(value=0.95),
+        gr.Checkbox(value=True),
+        gr.Slider(value=4),
+        gr.Slider(value=0.5),
     )
 
 
@@ -86,6 +107,55 @@ def gradio_app():
                     # value = "Quina és la finalitat del Servei Meteorològic de Catalunya?"
                 )
 
+                with gr.Row(variant="panel"):
+                    with gr.Accordion("Model parameters", open=False, visible=SHOW_MODEL_PARAMETERS_IN_UI):
+                        max_new_tokens = Slider(
+                            minimum=1,
+                            maximum=200,
+                            step=1,
+                            value=MAX_NEW_TOKENS,
+                            label="Max tokens"
+                        )
+                        repetition_penalty = Slider(
+                            minimum=0.1,
+                            maximum=10,
+                            step=0.1,
+                            value=1.2,
+                            label="Repetition penalty"
+                        )
+                        top_k = Slider(
+                            minimum=1,
+                            maximum=100,
+                            step=1,
+                            value=50,
+                            label="Top k"
+                        )
+                        top_p = Slider(
+                            minimum=0.01,
+                            maximum=0.99,
+                            value=0.95,
+                            label="Top p"
+                        )  
+                        do_sample = Checkbox(
+                            value=True, 
+                            label="Do sample"
+                        )
+                        num_beams = Slider(
+                            minimum=1, 
+                            maximum=8, 
+                            step=1,
+                            value=4,
+                            label="Beams"
+                        )
+                        temperature = Slider(
+                            minimum=0, 
+                            maximum=1,
+                            value=0.5,
+                            label="Temperature"
+                        )
+
+                        parameters_compontents = [max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature]
+
             with gr.Column(variant="panel"):
                 output = Textbox(
                     lines=11, label="Output", interactive=False, show_copy_button=True
@@ -96,6 +166,9 @@ def gradio_app():
                     )
                     submit_btn = Button("Submit", variant="primary", interactive=False)
 
+
+                
+
         input_.change(
             fn=change_interactive,
             inputs=[input_],
@@ -114,10 +187,18 @@ def gradio_app():
         )
 
         clear_btn.click(
-            fn=clear, inputs=[], outputs=[input_, output], queue=False, api_name=False
+            fn=clear, 
+            inputs=[], 
+            outputs=[input_, output] + parameters_compontents,
+              queue=False, 
+              api_name=False
         )
+        
         submit_btn.click(
-            fn=submit_input, inputs=[input_], outputs=[output], api_name="get-results"
+            fn=submit_input, 
+            inputs=[input_]+ parameters_compontents, 
+            outputs=[output], 
+            api_name="get-results"
         )
 
         with gr.Row():