change generation parameters values

app.py

@@ -7,7 +7,7 @@ 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))
+MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", default=200))
 SHOW_MODEL_PARAMETERS_IN_UI = os.environ.get("SHOW_MODEL_PARAMETERS_IN_UI", default="True") == "True"
 
 setup()
@@ -44,13 +44,13 @@ def submit_input(input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sa
 
 
     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
+        "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)
@@ -110,17 +110,17 @@ def gradio_app():
                 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,
+                            minimum=50,
+                            maximum=1000,
                             step=1,
                             value=MAX_NEW_TOKENS,
                             label="Max tokens"
                         )
                         repetition_penalty = Slider(
                             minimum=0.1,
-                            maximum=10,
+                            maximum=2.0,
                             step=0.1,
-                            value=1.2,
+                            value=1.0,
                             label="Repetition penalty"
                         )
                         top_k = Slider(
@@ -132,25 +132,25 @@ def gradio_app():
                         )
                         top_p = Slider(
                             minimum=0.01,
-                            maximum=0.99,
-                            value=0.95,
+                            maximum=1.0,
+                            value=1.0,
                             label="Top p"
                         )  
                         do_sample = Checkbox(
-                            value=True, 
+                            value=False, 
                             label="Do sample"
                         )
                         num_beams = Slider(
                             minimum=1, 
-                            maximum=8, 
+                            maximum=4, 
                             step=1,
-                            value=4,
+                            value=1,
                             label="Beams"
                         )
                         temperature = Slider(
-                            minimum=0, 
+                            minimum=0.1, 
                             maximum=1,
-                            value=0.5,
+                            value=0.35,
                             label="Temperature"
                         )
 

rag.py

@@ -38,7 +38,7 @@ class RAG:
 
         return context
         
-    def predict(self, instruction, context):
+    def predict(self, instruction, context, model_parameters):
 
         api_key = os.getenv("HF_TOKEN")
 
@@ -55,18 +55,18 @@ class RAG:
 
         payload = {
         "inputs": query,
-        "parameters": {"MAX_NEW_TOKENS": 1000, "TEMPERATURE": 0.25}
+        "parameters": model_parameters
         }
         
         response = requests.post(self.model_name, headers=headers, json=payload)
 
         return response.json()[0]["generated_text"].split("###")[-1][8:-1]
 
-    def get_response(self, prompt: str) -> str:
+    def get_response(self, prompt: str, model_parameters: dict) -> str:
         
         context = self.get_context(prompt)
 
-        response = self.predict(prompt, context)
+        response = self.predict(prompt, context, model_parameters)
 
         if not response:
             return self.NO_ANSWER_MESSAGE