Update app.py

app.py

@@ -8,6 +8,9 @@ import matplotlib.pyplot as plt
 import io
 import base64
 import os
+import huggingface_hub
+from huggingface_hub import hf_hub_download
+
 
 # Load label mapping
 label_to_int = pd.read_pickle('label_to_int.pkl')
@@ -47,8 +50,17 @@ def load_model():
     input_size = 768 + metadata_features
     log_reg = LogisticRegressionTorch(input_dim=input_size, output_dim=N_UNIQUE_CLASSES)
 
-    model_weights_path = os.getenv('MODEL_PATH')
-    weights = torch.load(model_weights_path, map_location=torch.device('cpu'))
+    #model_weights_path = os.getenv('MODEL_PATH')
+    #weights = torch.load(model_weights_path, map_location=torch.device('cpu'))
+
+
+    # Download the file
+    file_path = hf_hub_download(repo_id="mawairon/noo_test", filename="/gena-blastln-bs33-lr4e-05-S168.pth")
+")
+
+    # Now you can use the file_path to load your .pth file
+    weights = torch.load(file_path)
+
 
     base_model.load_state_dict(weights['model_state_dict'])
     log_reg.load_state_dict(weights['log_reg_state_dict'])