Update api/predict.py

api/predict.py

@@ -1,11 +1,8 @@
 import torch
 import numpy as np
 import re
-import os
 from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast
-
-scriptDir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-modelsDir = os.path.join(scriptDir, "models")
+from huggingface_hub import hf_hub_download
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 tokenizer = None
@@ -26,12 +23,20 @@ def load_resources():
     
     models = []
     for i in range(1, 6):
+        model_filename = f"ensemble_model_{i}.pth"
+        
+        print(f"downloading {model_filename}...")
+        model_path = hf_hub_download(
+            repo_id="codingcoolfun9ed/sentinelcheck-models",
+            filename=model_filename
+        )
+        
         model = DistilBertForSequenceClassification.from_pretrained(
             'distilbert-base-uncased',
             num_labels=num_classes,
             dropout=dropout
         )
-        model.load_state_dict(torch.load(os.path.join(modelsDir, f"ensemble_model_{i}.pth"), map_location=device))
+        model.load_state_dict(torch.load(model_path, map_location=device))
         model = model.to(device)
         model.eval()
         models.append(model)