Update streamlit app to download the classifier and vectorizer from the hub

app.py

@@ -10,6 +10,7 @@ import sys
 import os
 import json
 from urllib.parse import quote
+from huggingface_hub import hf_hub_download
 
 # Allow direct execution
 sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src'))  # noqa
@@ -32,6 +33,7 @@ st.set_page_config(
 
 MODEL_PATH = 'Xenova/sponsorblock-small_v2022.01.19'
 
+CLASSIFIER_PATH = 'Xenova/sponsorblock-classifier'
 
 @st.cache(allow_output_mutation=True)
 def persistdata():
@@ -54,6 +56,10 @@ def load_predict():
 
     tokenizer = AutoTokenizer.from_pretrained(evaluation_args.model_path)
 
+    # Save classifier and vectorizer
+    hf_hub_download(repo_id=CLASSIFIER_PATH, filename=classifier_args.classifier_file, cache_dir=classifier_args.classifier_dir)
+    hf_hub_download(repo_id=CLASSIFIER_PATH, filename=classifier_args.vectorizer_file, cache_dir=classifier_args.classifier_dir)
+
     def predict_function(video_id):
         if video_id not in predictions_cache:
             predictions_cache[video_id] = pred(