Update model.py

model.py

@@ -8,6 +8,8 @@ import sys
 import os
 import re
 from collections import Counter
+from huggingface_hub import hf_hub_download
+
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -190,7 +192,10 @@ def main():
 
     # Get the directory where this script is located
     script_dir = os.path.dirname(os.path.abspath(__file__))
-    CHECKPOINT_PATH = os.path.join(script_dir, "best_checkpoint.pth")
+    CHECKPOINT_PATH = hf_hub_download(
+    repo_id="VIKRAM989/image-label",
+    filename="best_checkpoint.pth"
+    )
     VOCAB_PATH = os.path.join(script_dir, "vocab.pkl")
 
     # load vocab