preload models

interfaces/cap.py

@@ -12,6 +12,7 @@ from huggingface_hub import HfApi
 from label_dicts import CAP_NUM_DICT, CAP_LABEL_NAMES
 
 from .utils import is_disk_full, release_model
+from ..utils import MODELS_PRELOADED, TOKENIZERS_PRELOADED
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -83,11 +84,12 @@ def build_huggingface_path(language: str, domain: str):
     else:
         return "poltextlab/xlm-roberta-large-pooled-cap"
 
-#@spaces.GPU
 def predict(text, model_id, tokenizer_id):
     device = torch.device("cpu")
-    model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN).to(device)
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+    
+    print(MODELS_PRELOADED, TOKENIZERS_PRELOADED)
+    model = MODELS_PRELOADED[model_id].to(device)
+    tokenizer = TOKENIZERS_PRELOADED[tokenizer_id]
 
     inputs = tokenizer(text,
                        max_length=256,

utils.py

@@ -25,6 +25,9 @@ from interfaces.ontolisst import build_huggingface_path as hf_ontolisst_path
 
 from huggingface_hub import scan_cache_dir
 
+MODELS_PRELOADED = []
+TOKENIZERS_PRELOADED = []
+
 HF_TOKEN = os.environ["hf_read"]
 
 # should be a temporary solution
@@ -54,9 +57,9 @@ tokenizers = ["xlm-roberta-large"]
 
 def download_hf_models():
     for model_id in models:
-        AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
+        MODELS_PRELOADED[model_id] = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
     for tokenizer_id in tokenizers:
-        AutoTokenizer.from_pretrained(tokenizer_id)
+        TOKENIZERS_PRELOADED[tokenizer_id] = AutoTokenizer.from_pretrained(tokenizer_id)
         
         
 def df_h():