secrets

app/download_models.py

@@ -1,6 +1,9 @@
 import os
 from huggingface_hub import hf_hub_download
 
+# Get the Hugging Face token from environment variable (automatically provided from the secret)
+HUGGINGFACE_TOKEN = os.getenv("key")
+
 
 def create_directory(dir_path):
     """
@@ -13,15 +16,17 @@ def create_directory(dir_path):
         print(f"Directory {dir_path} already exists.")
 
 
-def download_model(repo_id, filename, destination_dir):
+def download_model(repo_id, filename, destination_dir, token):
     """
     Download the model from Hugging Face and move it to the destination directory.
     """
     # Create destination directory if it doesn't exist
     create_directory(destination_dir)
 
-    # Download the file from the Hugging Face Hub
-    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    # Download the file from the Hugging Face Hub using the token
+    model_path = hf_hub_download(
+        repo_id=repo_id, filename=filename, use_auth_token=token
+    )
 
     # Move the downloaded file to the destination directory
     destination_path = os.path.join(destination_dir, filename)
@@ -47,11 +52,16 @@ if __name__ == "__main__":
     oracle_cards_graph_dir = "./data/oracle_cards_graph"
 
     # Download the card_graph.pkl from uripper/magic-gnn
-    card_graph_path = download_model(magic_gnn_repo, card_graph_file, card_graph_dir)
+    card_graph_path = download_model(
+        magic_gnn_repo, card_graph_file, card_graph_dir, HUGGINGFACE_TOKEN
+    )
 
     # Download the oracle_cards_graph.pkl from uripper/oracle_cards_graph
     oracle_cards_graph_path = download_model(
-        oracle_cards_graph_repo, oracle_cards_graph_file, oracle_cards_graph_dir
+        oracle_cards_graph_repo,
+        oracle_cards_graph_file,
+        oracle_cards_graph_dir,
+        HUGGINGFACE_TOKEN,
     )
 
     # Print final message with full paths