cache models during build

app.py

@@ -1,3 +1,10 @@
+import os
+PATH = '/data/' # at least 150GB storage needs to be attached
+os.environ['TRANSFORMERS_CACHE'] = PATH
+os.environ['HF_HOME'] = PATH
+os.environ['HF_DATASETS_CACHE'] = PATH
+os.environ['TORCH_HOME'] = PATH
+
 import gradio as gr
 
 from interfaces.cap import demo as cap_demo
@@ -6,7 +13,7 @@ from interfaces.sentiment import demo as sentiment_demo
 from interfaces.emotion import demo as emotion_demo
 from interfaces.ner import demo as ner_demo
 from interfaces.ner import download_models as download_spacy_models
-
+from utils import download_hf_models
 
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -29,6 +36,7 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
+    download_hf_models()
     download_spacy_models()
     demo.launch()
 

utils.py

@@ -0,0 +1,28 @@
+import os
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+from interfaces.cap import languages as languages_cap
+from interfaces.cap import domains as domains_cap
+
+from interfaces.cap import build_huggingface_path as hf_cap_path
+from interfaces.manifesto import build_huggingface_path as hf_manifesto_path
+from interfaces.sentiment import build_huggingface_path as hf_sentiment_path
+from interfaces.emotion import build_huggingface_path as hf_emotion_path
+
+HF_TOKEN = os.environ["hf_read"]
+
+# should be a temporary solution
+models = [hf_manifesto_path(""), hf_sentiment_path(""), hf_emotion_path("")]
+for language in languages_cap:
+    for domain in domains_cap:
+        models.append(hf_cap_path(language, domain))
+
+tokenizers = ["xlm-roberta-large"]
+
+def download_hf_models():
+    for model_id in models:
+        AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto",
+                                                                   token=HF_TOKEN)
+    for tokenizer_id in tokenizers:
+        AutoTokenizer.from_pretrained(tokenizer_id)
+