Update train.py

train.py

@@ -1,31 +1,13 @@
 # ==========================================
-# 0. BRUTE FORCE SETUP & AUTHENTICATION
+# 0. PURE PYTHON SETUP & AUTO-FETCHER
 # ==========================================
 import os
 import sys
-import urllib.request
-import subprocess
+from huggingface_hub import hf_hub_download, login
 
 # Disable UI prompts at the deepest system level
 os.environ["HF_HUB_DISABLE_INTERACTIVE_PROMPTS"] = "1"
-os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" # Sometimes transfer agents cause prompt loops
-
-HF_USERNAME = "Proteinrequired" 
-SPACE_NAME = "enterprise-email-triage"
-
-print("Checking for required environment files...")
-# THE FIX: Added openenv.yaml to this list!
-files_to_download = ["env.py", "dataset.json", "reward_system.py", "openenv.yaml"]
-base_url = f"https://huggingface.co/spaces/{HF_USERNAME}/{SPACE_NAME}/resolve/main/"
-
-for file in files_to_download:
-    if not os.path.exists(file):
-        print(f"Downloading {file} from repository...")
-        try:
-            urllib.request.urlretrieve(base_url + file, file)
-            print(f"✅ {file} downloaded successfully.")
-        except Exception as e:
-            print(f"❌ Failed to download {file}")
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 
 raw_token = os.environ.get("HF_TOKEN")
 if not raw_token:
@@ -35,14 +17,29 @@ if not raw_token:
 hf_token = raw_token.strip()
 os.environ["HF_TOKEN"] = hf_token
 
-print("🔑 Forcing System-Level Hugging Face Login...")
-# This writes the token directly to the machine's hard drive so no library can ignore it
-try:
-    subprocess.run(["huggingface-cli", "login", "--token", hf_token, "--add-to-git-credential"], check=True)
-    print("✅ System Login Successful!")
-except subprocess.CalledProcessError:
-    print("❌ System Login Failed. Check if your token is correct.")
-    sys.exit(1)
+print("🔑 Authenticating via pure Python...")
+# Python-level login instead of system-level CLI
+login(token=hf_token, add_to_git_credential=True)
+
+HF_USERNAME = "Proteinrequired" 
+SPACE_NAME = "enterprise-email-triage"
+files_to_download = ["env.py", "dataset.json", "reward_system.py", "openenv.yaml"]
+
+print(f"Fetching master environment files from {SPACE_NAME}...")
+for file in files_to_download:
+    try:
+        # hf_hub_download automatically handles private spaces!
+        hf_hub_download(
+            repo_id=f"{HF_USERNAME}/{SPACE_NAME}", 
+            repo_type="space", 
+            filename=file,
+            token=hf_token, # Passes the token securely
+            local_dir="." # Saves right next to train.py
+        )
+        print(f"✅ {file} downloaded successfully.")
+    except Exception as e:
+        print(f"❌ Failed to download {file}. Error: {e}")
+        sys.exit(1)
 # ==========================================
 # 1. THE STABLE ENVIRONMENT INSTALL
 # ==========================================