utils.py

interfaces/cap.py

@@ -12,7 +12,6 @@ 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"]
 
@@ -86,10 +85,8 @@ def build_huggingface_path(language: str, domain: str):
 
 def predict(text, model_id, tokenizer_id):
     device = torch.device("cpu")
-    
-    print(MODELS_PRELOADED, TOKENIZERS_PRELOADED)
-    model = MODELS_PRELOADED[model_id].to(device)
-    tokenizer = TOKENIZERS_PRELOADED[tokenizer_id]
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     inputs = tokenizer(text,
                        max_length=256,