Update app.py

app.py

@@ -30,7 +30,7 @@ class LogQueueCallback(TrainerCallback):
             log_str = f"Step {state.global_step}: {json.dumps(logs)}\n"
             self.log_queue.put(log_str)
 
-def get_user_info(token):
+def get_username(token):
     """Retrieves the username from the HF token."""
     if not token:
         return None
@@ -59,22 +59,20 @@ def train_thread_target(
     result_queue
 ):
     """
-    Background thread for training.
+    Background thread for training and pushing to user profile.
     """
     try:
-        # 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)
+        # 0. Auth & Identity
+        final_token = token or os.environ.get("HF_TOKEN")
+        username = get_username(final_token)
+        
         if not username:
-            raise ValueError("Invalid Hugging Face Token. Please provide a token in the UI or set HF_TOKEN in your environment secrets.")
+            raise ValueError("Invalid or missing Hugging Face Token. Ensure the token is provided or set as HF_TOKEN secret.")
             
-        # 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")
+        # Target path is now the USER'S profile
+        full_repo_id = f"{username}/{model_name}"
+        log_queue.put(f"🚀 Initializing for user: {username}\n")
+        log_queue.put(f"📦 Target Repository: https://huggingface.co/{full_repo_id}\n")
 
         # Validation for Transformer logic
         if n_embd % n_head != 0:
@@ -83,6 +81,7 @@ def train_thread_target(
         # 1. Load Dataset
         log_queue.put(f"📚 Loading dataset: {dataset_id} (Limit: {sample_limit})...\n")
         try:
+            # We use the train split; user can specify limit
             dataset = load_dataset(dataset_id, split=f"train[:{int(sample_limit)}]")
         except Exception as e:
             raise ValueError(f"Error loading dataset: {e}")
@@ -128,7 +127,7 @@ def train_thread_target(
         log_queue.put("🏋️ Starting Training Loop...\n")
         
         training_args = TrainingArguments(
-            output_dir="./results",
+            output_dir="./local_results",
             overwrite_output_dir=True,
             num_train_epochs=epochs,
             per_device_train_batch_size=int(batch_size),
@@ -137,7 +136,7 @@ def train_thread_target(
             weight_decay=weight_decay,
             warmup_steps=int(warmup_steps),
             logging_steps=10,
-            save_strategy="no", # Save only at the end
+            save_strategy="no", 
             push_to_hub=False,
             report_to="none",
             use_cpu=not torch.cuda.is_available(),
@@ -154,18 +153,18 @@ def train_thread_target(
 
         trainer.train()
 
-        # 5. Push to Hub
-        log_queue.put(f"☁️ Pushing weights to https://huggingface.co/{full_repo_id}...\n")
+        # 5. Push to User's Personal Hub
+        log_queue.put(f"☁️ Uploading model to your profile...\n")
         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}")
+        result_queue.put(f"🎉 Success! Published to: https://huggingface.co/{full_repo_id}")
     
     except Exception as e:
         log_queue.put(f"❌ Error: {str(e)}\n")
         result_queue.put(None)
 
-# --- Main Generator Function ---
+# --- Generator for UI updates ---
 
 def train_and_push_generator(
     token, dataset_id, model_name, 
@@ -173,11 +172,10 @@ def train_and_push_generator(
     epochs, lr, weight_decay, warmup_steps,
     batch_size, grad_accumulation, sample_limit
 ):
-    # 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.", ""
+        yield "Error: No Hugging Face Token found. Please enter a 'Write' token below.", ""
         return
     
     log_queue = queue.Queue()
@@ -204,31 +202,31 @@ def train_and_push_generator(
         
     if not result_queue.empty():
         result = result_queue.get()
-        yield logs_history, result or "Failed. Check logs for errors."
+        yield logs_history, result or "Training failed. See logs."
     else:
-        yield logs_history, "Process finished unexpectedly."
+        yield logs_history, "Process interrupted."
 
 # --- UI Layout ---
 
-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 the **Auto-PreTrain** organization.")
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate")) as demo:
+    gr.Markdown("# 🚀 Personal Auto-PreTrain")
+    gr.Markdown("Configure a custom GPT-2 architecture and train it directly to **your personal** Hugging Face profile.")
     
     with gr.Row():
         hf_token = gr.Textbox(
-            label="Hugging Face Write Token (Optional if set in Secrets)", 
+            label="HF Write Token", 
             placeholder="hf_...", 
             type="password",
-            info="Leave blank if you have 'HF_TOKEN' set in your environment/secrets. Token must have write access to 'Auto-PreTrain' org."
+            info="Required to create the repo on your profile. Must have 'Write' permissions."
         )
         model_name_input = gr.Textbox(
-            label="Model Repository Name", 
-            value="my-tiny-gpt2",
-            placeholder="e.g. tiny-coder-v1"
+            label="Model Name", 
+            value="my-custom-gpt2",
+            placeholder="e.g. tiny-stories-v1"
         )
 
     with gr.Tabs():
-        with gr.TabItem("1. Dataset & Data"):
+        with gr.TabItem("1. Data Selection"):
             with gr.Row():
                 dataset_input = gr.Textbox(
                     label="Dataset ID", 
@@ -236,40 +234,39 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="orange", secondary_hue="gray
                     placeholder="e.g. wikitext"
                 )
                 sample_limit = gr.Number(
-                    label="Sample Limit", 
-                    value=1000, 
-                    precision=0,
-                    info="Number of rows to use for training"
+                    label="Training Samples", 
+                    value=500, 
+                    precision=0
                 )
             context_length = gr.Slider(
                 minimum=64, maximum=1024, value=128, step=64, 
-                label="Max Context Length (Sequence Length)"
+                label="Max Context Length"
             )
 
-        with gr.TabItem("2. Model Architecture"):
+        with gr.TabItem("2. Architecture"):
             with gr.Row():
-                layers = gr.Slider(minimum=1, maximum=24, value=4, step=1, label="Number of Layers")
-                embd = gr.Slider(minimum=64, maximum=1024, value=256, step=64, label="Embedding Dimension")
+                layers = gr.Slider(minimum=1, maximum=12, value=2, step=1, label="Layers")
+                embd = gr.Slider(minimum=64, maximum=1024, value=128, step=64, label="Embedding Dim")
             with gr.Row():
-                heads = gr.Slider(minimum=2, maximum=16, value=8, step=2, label="Attention Heads")
-                gr.Markdown("Note: Embedding dimension must be divisible by attention heads.")
+                heads = gr.Slider(minimum=2, maximum=16, value=4, step=2, label="Attention Heads")
+                gr.Markdown("_Note: Embedding Dim must be divisible by Attention Heads._")
 
-        with gr.TabItem("3. Training Hyperparameters"):
+        with gr.TabItem("3. Training Settings"):
             with gr.Row():
-                epochs = gr.Slider(minimum=1, maximum=50, value=1, step=1, label="Epochs")
+                epochs = gr.Slider(minimum=1, maximum=20, value=1, step=1, label="Epochs")
                 lr = gr.Number(label="Learning Rate", value=5e-4)
             with gr.Row():
-                batch_size = gr.Slider(minimum=1, maximum=64, value=8, step=1, label="Batch Size (per device)")
-                grad_accumulation = gr.Slider(minimum=1, maximum=32, value=1, step=1, label="Gradient Accumulation Steps")
+                batch_size = gr.Slider(minimum=1, maximum=32, value=4, step=1, label="Batch Size")
+                grad_accumulation = gr.Slider(minimum=1, maximum=16, value=1, step=1, label="Grad Accumulation")
             with gr.Row():
                 weight_decay = gr.Slider(minimum=0.0, maximum=0.1, value=0.01, step=0.01, label="Weight Decay")
-                warmup_steps = gr.Number(label="Warmup Steps", value=100, precision=0)
+                warmup_steps = gr.Number(label="Warmup Steps", value=50, precision=0)
 
-    train_btn = gr.Button("🚀 Start Pre-Training", variant="primary")
+    train_btn = gr.Button("🔥 Start Training & Push to My Profile", variant="primary")
     
     with gr.Row():
-        log_output = gr.Code(label="Live Training Logs", language="json", lines=15)
-        status_output = gr.Textbox(label="Status & Hub Link", interactive=False)
+        log_output = gr.Code(label="Training Progress", language="json", lines=12)
+        status_output = gr.Textbox(label="Final Status", interactive=False)
 
     train_btn.click(
         fn=train_and_push_generator,