added a checkbox that can disable the web search

app.py

@@ -40,6 +40,12 @@ with st.form("request_form"):
         max_chars=2048,
     )
 
+    web_search: bool = st.checkbox(
+        label="Web search",
+        value=True,
+        help="If checked, the model will get your prompt as well as some web search results."
+    )
+
     submitted: bool = st.form_submit_button(
         label="Generate",
         help="Generate the output text.",
@@ -48,7 +54,12 @@ with st.form("request_form"):
 
     if submitted:
         try:
-            output = on_form_submit(selected_model_name, output_length, submitted_prompt)
+            output = on_form_submit(
+                selected_model_name,
+                output_length,
+                submitted_prompt,
+                web_search,
+            )
         except CudaError as e:
             st.error("Out of memory. Please try a smaller model, shorter prompt, or a smaller output length.")
         except (ValueError, TypeError, RuntimeError) as e:

hanlde_form_submit.py

@@ -52,15 +52,20 @@ def generate_text(
         pipeline: GroupedSamplingPipeLine,
         prompt: str,
         output_length: int,
+        web_search: bool,
 ) -> str:
     """
     Generates text using the given pipeline.
     :param pipeline: The pipeline to use. GroupedSamplingPipeLine.
     :param prompt: The prompt to use. str.
     :param output_length: The size of the text to generate in tokens. int > 0.
+    :param web_search: Whether to use web search or not. bool.
     :return: The generated text. str.
     """
-    better_prompt = rewrite_prompt(prompt)
+    if web_search:
+        better_prompt = rewrite_prompt(prompt)
+    else:
+        better_prompt = prompt
     return pipeline(
         prompt_s=better_prompt,
         max_new_tokens=output_length,
@@ -69,12 +74,18 @@ def generate_text(
     )["generated_text"]
 
 
-def on_form_submit(model_name: str, output_length: int, prompt: str) -> str:
+def on_form_submit(
+        model_name: str,
+        output_length: int,
+        prompt: str,
+        web_search: bool
+) -> str:
     """
     Called when the user submits the form.
     :param model_name: The name of the model to use.
     :param output_length: The size of the groups to use.
     :param prompt: The prompt to use.
+    :param web_search: Whether to use web search or not.
     :return: The output of the model.
     :raises ValueError: If the model name is not supported, the output length is <= 0,
      the prompt is empty or longer than
@@ -99,6 +110,7 @@ def on_form_submit(model_name: str, output_length: int, prompt: str) -> str:
         pipeline=pipeline,
         prompt=prompt,
         output_length=output_length,
+        web_search=web_search,
     )
     generation_end_time = time()
     generation_time = generation_end_time - generation_start_time