Update requirements and make necessary code changes

.ipynb_checkpoints/model_loader-checkpoint.py

@@ -1,3 +1,5 @@
+from huggingface_hub import hf_hub_download
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -628,14 +630,11 @@ def load_model():
         model, tokenizer = load_T5_model_classification(checkpoint, num_labels, mixed, full, deepspeed)
     
 
-    checkpoint_dir = model.config.name_or_path  # This will point to the local directory
-
-    print(checkpoint_dir)
-    # Construct the path to the custom checkpoint file
-    best_model_path = os.path.join(checkpoint_dir, 'cpt.pth')
+    # Download the file
+    local_file = hf_hub_download(repo_id=checkpoint, filename="cpt.pth")
 
     # Load the best model state
-    state_dict = torch.load(best_model_path, weights_only=True)
+    state_dict = torch.load(local_file, weights_only=True)
     model.load_state_dict(state_dict)
 
     return model, tokenizer

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -7,3 +7,4 @@ pandas>=1.1.0
 numpy>=1.19.0
 scikit-learn>=0.24.0
 sentencepiece
+huggingface_hub>=0.15.0

model_loader.py

@@ -1,3 +1,5 @@
+from huggingface_hub import hf_hub_download
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -628,14 +630,11 @@ def load_model():
         model, tokenizer = load_T5_model_classification(checkpoint, num_labels, mixed, full, deepspeed)
     
 
-    checkpoint_dir = model.config.name_or_path  # This will point to the local directory
-
-    print(checkpoint_dir)
-    # Construct the path to the custom checkpoint file
-    best_model_path = os.path.join(checkpoint_dir, 'cpt.pth')
+    # Download the file
+    local_file = hf_hub_download(repo_id=checkpoint, filename="cpt.pth")
 
     # Load the best model state
-    state_dict = torch.load(best_model_path, weights_only=True)
+    state_dict = torch.load(local_file, weights_only=True)
     model.load_state_dict(state_dict)
 
     return model, tokenizer

requirements.txt

@@ -7,3 +7,4 @@ pandas>=1.1.0
 numpy>=1.19.0
 scikit-learn>=0.24.0
 sentencepiece
+huggingface_hub>=0.15.0