downloading models at the start of the app and not at usage time

app.py

@@ -4,8 +4,10 @@ In the demo, the user can write a prompt
  and the model will generate a response using the grouped sampling algorithm.
 """
 import os
+from time import time
 
 import streamlit as st
+from grouped_sampling import GroupedSamplingPipeLine
 from torch.cuda import CudaError
 from huggingface_hub import logging as hf_hub_logging
 
@@ -13,6 +15,27 @@ from available_models import AVAILABLE_MODELS
 from hanlde_form_submit import on_form_submit
 
 
+def create_pipeline(model_name: str, group_size: int) -> 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.
+    """
+    st.write(f"Starts creating pipeline with model: {model_name}")
+    pipeline_start_time = time()
+    pipeline = GroupedSamplingPipeLine(
+        model_name=model_name,
+        group_size=group_size,
+        end_of_sentence_stop=False,
+        top_k=50,
+    )
+    pipeline_end_time = time()
+    pipeline_time = pipeline_end_time - pipeline_start_time
+    st.write(f"Finished creating pipeline with model: {model_name} in {pipeline_time:,.2f} seconds.")
+    return pipeline
+
+
 hf_hub_logging.set_verbosity_error()
 
 st.set_page_config(
@@ -20,6 +43,10 @@ st.set_page_config(
     layout="wide",
 )
 
+pipelines = {
+    model_name: create_pipeline(model_name, 1024) for model_name in AVAILABLE_MODELS[1:]
+}
+
 with st.form("request_form"):
     selected_model_name: str = st.selectbox(
         label="בחרו מודל",
@@ -50,7 +77,7 @@ with st.form("request_form"):
     if submitted:
         try:
             output = on_form_submit(
-                selected_model_name,
+                pipelines[selected_model_name],
                 output_length,
                 submitted_prompt,
             )

hanlde_form_submit.py

@@ -1,28 +1,7 @@
 from time import time
 
 import streamlit as st
-from grouped_sampling import GroupedSamplingPipeLine, is_supported, UnsupportedModelNameException
-
-
-def create_pipeline(model_name: str, group_size: int) -> 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.
-    """
-    st.write(f"Starts creating pipeline with model: {model_name}")
-    pipeline_start_time = time()
-    pipeline = GroupedSamplingPipeLine(
-        model_name=model_name,
-        group_size=group_size,
-        end_of_sentence_stop=False,
-        top_k=50,
-    )
-    pipeline_end_time = time()
-    pipeline_time = pipeline_end_time - pipeline_start_time
-    st.write(f"Finished creating pipeline with model: {model_name} in {pipeline_time:,.2f} seconds.")
-    return pipeline
+from grouped_sampling import GroupedSamplingPipeLine
 
 
 def generate_text(
@@ -46,13 +25,13 @@ def generate_text(
 
 
 def on_form_submit(
-        model_name: str,
+        pipeline: GroupedSamplingPipeLine,
         output_length: int,
         prompt: str,
 ) -> str:
     """
     Called when the user submits the form.
-    :param model_name: The name of the model to use.
+    :param pipeline: The pipeline to use. GroupedSamplingPipeLine.
     :param output_length: The size of the groups to use.
     :param prompt: The prompt to use.
     :return: The output of the model.
@@ -64,19 +43,6 @@ def on_form_submit(
     """
     if len(prompt) == 0:
         raise ValueError("The prompt must not be empty.")
-    if not is_supported(model_name):
-        raise UnsupportedModelNameException(model_name)
-    st.write(f"Loading model: {model_name}...")
-    print(f"Loading model: {model_name}...")
-    loading_start_time = time()
-    pipeline = create_pipeline(
-        model_name=model_name,
-        group_size=output_length,
-    )
-    loading_end_time = time()
-    loading_time = loading_end_time - loading_start_time
-    st.write(f"Finished loading model: {model_name}  in {loading_time:,.2f} seconds.")
-    print(f"Finished loading model: {model_name}  in {loading_time:,} seconds.")
     st.write("Generating text...")
     print("Generating text...")
     generation_start_time = time()
@@ -89,8 +55,4 @@ def on_form_submit(
     generation_time = generation_end_time - generation_start_time
     st.write(f"Finished generating text in {generation_time:,.2f} seconds.")
     print(f"Finished generating text in {generation_time:,.2f} seconds.")
-    if not isinstance(generated_text, str):
-        raise RuntimeError(f"The model {model_name} did not generate any text.")
-    if len(generated_text) == 0:
-        raise RuntimeError(f"The model {model_name} did not generate any text.")
     return generated_text