moved functions to new file

app.py

@@ -1,41 +1,24 @@
 """
 The Streamlit app for the project demo.
-In the demo, the user can write a prompt and the model will generate a response using the grouped sampling algorithm.
+In the demo, the user can write a prompt
+ and the model will generate a response using the grouped sampling algorithm.
 """
 
 import streamlit as st
-from grouped_sampling import GroupedSamplingPipeLine
 
-available_models_list = "https://huggingface.co/models?pipeline_tag=text-generation&sort=downloads"
+from hanlde_form_submit import on_form_submit
 
 
-def create_pipeline(model_name: str, group_size) -> GroupedSamplingPipeLine:
-    """
-    Creates a pipeline with the given model name and group size.
-    :param model_name: The name of the model to use.
-    :param group_size: The size of the groups to use.
-    :return: A pipeline with the given model name and group size.
-    """
-    return GroupedSamplingPipeLine(model_name=model_name, group_size=group_size)
-
-
-def on_form_submit(model_name: str, group_size: int, prompt: str) -> str:
-    """
-    Called when the user submits the form.
-    :param model_name: The name of the model to use.
-    :param group_size: The size of the groups to use.
-    :param prompt: The prompt to use.
-    :return: The output of the model.
-    """
-    pipeline = create_pipeline(model_name, group_size)
-    return pipeline(prompt)["generated_text"]
+AVAILABLE_MODEL_NAMES = "https://huggingface.co/models?pipeline_tag=text-generation&sort=downloads"
 
 
 with st.form("request_form"):
     selected_model_name: str = st.text_input(
         label="Model name",
         value="gpt2",
-        help=f"The name of the model to use. Must be a model from this list: {available_models_list}"
+        help=f"The name of the model to use."
+             f" Must be a model from this list:"
+             f" {AVAILABLE_MODEL_NAMES}"
     )
 
     output_length: int = st.number_input(

hanlde_form_submit.py

@@ -0,0 +1,30 @@
+from grouped_sampling import GroupedSamplingPipeLine
+
+
+def create_pipeline(model_name: str, group_size) -> GroupedSamplingPipeLine:
+    """
+    Creates a pipeline with the given model name and group size.
+    :param model_name: The name of the model to use.
+    :param group_size: The size of the groups to use.
+    :return: A pipeline with the given model name and group size.
+    """
+    return GroupedSamplingPipeLine(
+        model_name=model_name,
+        group_size=group_size,
+        end_of_sentence_stop=True,
+    )
+
+
+def on_form_submit(model_name: str, group_size: int, prompt: str) -> str:
+    """
+    Called when the user submits the form.
+    :param model_name: The name of the model to use.
+    :param group_size: The size of the groups to use.
+    :param prompt: The prompt to use.
+    :return: The output of the model.
+    """
+    pipeline = create_pipeline(
+        model_name,
+        group_size,
+    )
+    return pipeline(prompt)["generated_text"]