Update app.py

app.py

@@ -62,11 +62,18 @@ def train_thread_target(
     Background thread for training.
     """
     try:
-        username = get_user_info(token)
+        # Check if the token is provided or should be pulled from Environment Secrets
+        final_token = token
+        if not final_token:
+            final_token = os.environ.get("HF_TOKEN")
+            
+        username = get_user_info(final_token)
         if not username:
-            raise ValueError("Invalid Hugging Face Token. Could not authenticate.")
+            raise ValueError("Invalid Hugging Face Token. Please provide a token in the UI or set HF_TOKEN in your environment secrets.")
             
-        full_repo_id = f"{username}/{model_name}"
+        # Updated to push specifically to the Auto-PreTrain organization
+        org_name = "Auto-PreTrain"
+        full_repo_id = f"{org_name}/{model_name}"
         log_queue.put(f"🚀 Initializing for {full_repo_id}...\n")
 
         # 1. Load Dataset
@@ -145,8 +152,8 @@ def train_thread_target(
 
         # 5. Push to Hub
         log_queue.put(f"☁️ Pushing weights to https://huggingface.co/{full_repo_id}...\n")
-        model.push_to_hub(full_repo_id, token=token)
-        tokenizer.push_to_hub(full_repo_id, token=token)
+        model.push_to_hub(full_repo_id, token=final_token)
+        tokenizer.push_to_hub(full_repo_id, token=final_token)
 
         result_queue.put(f"🎉 Success! Model published to: https://huggingface.co/{full_repo_id}")
     
@@ -162,15 +169,18 @@ def train_and_push_generator(
     epochs, lr, weight_decay, warmup_steps,
     batch_size, grad_accumulation, sample_limit
 ):
-    if not token:
-        yield "Error: Hugging Face Token is required.", ""
+    # If UI token is empty, we attempt to use the environment variable secret
+    effective_token = token or os.environ.get("HF_TOKEN")
+
+    if not effective_token:
+        yield "Error: No Hugging Face Token found. Either enter it below or set HF_TOKEN in Secrets.", ""
         return
     
     log_queue = queue.Queue()
     result_queue = queue.Queue()
     
     t = threading.Thread(target=train_thread_target, args=(
-        token, dataset_id, model_name, 
+        effective_token, dataset_id, model_name, 
         num_layers, n_embd, n_head, context_length,
         epochs, lr, weight_decay, warmup_steps,
         batch_size, grad_accumulation, sample_limit,
@@ -198,14 +208,14 @@ def train_and_push_generator(
 
 with gr.Blocks(theme=gr.themes.Default(primary_hue="orange", secondary_hue="gray")) as demo:
     gr.Markdown("# 🔥 Advanced Auto-PreTrain")
-    gr.Markdown("Configure your transformer architecture and train it directly to your Hugging Face account.")
+    gr.Markdown("Configure your transformer architecture and train it directly to the **Auto-PreTrain** organization.")
     
     with gr.Row():
         hf_token = gr.Textbox(
-            label="Hugging Face Write Token", 
+            label="Hugging Face Write Token (Optional if set in Secrets)", 
             placeholder="hf_...", 
             type="password",
-            info="Get your token at huggingface.co/settings/tokens (must have 'Write' access)"
+            info="Leave blank if you have 'HF_TOKEN' set in your environment/secrets. Token must have write access to 'Auto-PreTrain' org."
         )
         model_name_input = gr.Textbox(
             label="Model Repository Name",