added a function to download models from the hub

download_repo.py

@@ -0,0 +1,26 @@
+import urllib3
+
+from huggingface_hub import snapshot_download
+
+
+def change_default_timeout(new_timeout: int) -> None:
+    """
+    Changes the default timeout for downloading repositories from the Hugging Face Hub.
+    Prevents the following errors:
+    urllib3.exceptions.ReadTimeoutError: HTTPSConnectionPool(host='huggingface.co', port=443):
+    Read timed out. (read timeout=10)
+    """
+    urllib3.util.timeout.DEFAULT_TIMEOUT = new_timeout
+
+
+def download_repository(name: str) -> None:
+    """Downloads a repository from the Hugging Face Hub."""
+    number_of_seconds_in_a_day: int = 86_400
+    change_default_timeout(number_of_seconds_in_a_day)
+    snapshot_download(
+        repo_id=name,
+        etag_timeout=number_of_seconds_in_a_day,
+        resume_download=True,
+        repo_type="model",
+        library_name="pytorch"
+    )

hanlde_form_submit.py

@@ -3,6 +3,7 @@ from time import time
 import streamlit as st
 from grouped_sampling import GroupedSamplingPipeLine
 
+from download_repo import download_repository
 from prompt_engeneering import rewrite_prompt
 from supported_models import get_supported_model_names
 
@@ -18,6 +19,14 @@ def create_pipeline(model_name: str, group_size: int) -> GroupedSamplingPipeLine
     :return: A pipeline with the given model name and group size.
     """
     print(f"Starts downloading model: {model_name} from the internet.")
+    download_repository_start_time = time()
+    st.write(f"Starts downloading model: {model_name} from the internet.")
+    download_repository(model_name)
+    download_repository_end_time = time()
+    download_time = download_repository_end_time - download_repository_start_time
+    st.write(f"Finished downloading model: {model_name} from the internet in {download_time} seconds.")
+    st.write(f"Starts creating pipeline with model: {model_name}")
+    pipeline_start_time = time()
     pipeline = GroupedSamplingPipeLine(
         model_name=model_name,
         group_size=group_size,
@@ -25,7 +34,9 @@ def create_pipeline(model_name: str, group_size: int) -> GroupedSamplingPipeLine
         temp=0.5,
         top_p=0.6,
     )
-    print(f"Finished downloading model: {model_name} from the internet.")
+    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} seconds.")
     return pipeline
 
 

on_server_start.py

@@ -1,20 +1,7 @@
 """
 A script that is run when the server starts.
 """
-from huggingface_hub import snapshot_download
-
-
-def download_model(model_name: str):
-    """
-    Downloads a model from hugging face hub to the disk but not to the RAM.
-    :param model_name: The name of the model to download.
-    """
-    number_of_seconds_in_a_day: int = 86_400
-    snapshot_download(
-        repo_id=model_name,
-        etag_timeout=number_of_seconds_in_a_day,
-        resume_download=True,
-    )
+from download_repo import download_repository
 
 
 def download_useful_models():
@@ -27,7 +14,7 @@ def download_useful_models():
         "facebook/opt-125m",
     )
     for model_name in useful_models:
-        download_model(model_name)
+        download_repository(model_name)
 
 
 def main():