Update train.py

train.py

@@ -25,7 +25,7 @@ dst_path = "./creditcard.Rdata"
 wget.download(url, dst_path)
 
 # Define the directory to save the model (same as dataset directory)
-dataset_dir = os.path.dirname(dst_path)  # Make sure this is defined before saving the model
+dataset_dir = os.path.dirname(dst_path)
 
 # Load the dataset
 parsed_res = rdata.parser.parse_file(dst_path)
@@ -82,16 +82,21 @@ joblib.dump(model_pipeline, saved_model_path)
 
 print(f"Model Serialized and Saved to {saved_model_path}")
 
-# Login to Hugging Face using the API token
+# Get the Hugging Face API token securely from the secret environment variable
+api_token = os.getenv('FraudDemoWrite')  # Access the secret directly using its name
+
+if not api_token:
+    raise ValueError("Hugging Face API token not found. Ensure 'FraudDemoWrite' is set as a secret in your environment.")
+
+# Initialize Hugging Face API
 api = HfApi()
-api.set_access_token(FraudDemoWrite)
 
 # Create or use an existing repository on Hugging Face Hub
 repo_name = "kgauvin603/creditcard-fraud-detection"  # Replace with your desired repo name
-repo_url = api.create_repo(repo_id=repo_name, exist_ok=True)
+repo_url = api.create_repo(repo_id=repo_name, token=api_token, exist_ok=True)
 
 # Initialize the repository
-repo = Repository(local_dir="hf_model_repo", clone_from=repo_url)
+repo = Repository(local_dir="hf_model_repo", clone_from=repo_url, use_auth_token=api_token)
 
 # Move the model to the repository folder
 os.rename(saved_model_path, os.path.join("hf_model_repo", "model.joblib"))