Update app.py

app.py

@@ -3,20 +3,30 @@ import subprocess
 import os
 import sys
 from datetime import datetime
+from transformers import AutoTokenizer
+from peft import AutoPeftModelForQuestionAnswering
+from huggingface_hub import login
 
 def run_training(model_name):
     """Execute the training script and save to Hub"""
     if not model_name.strip():
         return "❌ Please enter a model name!"
     
-    # Set environment variable for the script
+    # Set environment variables for the script
     os.environ['MODEL_NAME'] = model_name.strip()
     
+    # Get HF token from environment (set as Space secret)
+    hf_token = os.environ.get('roberta_token')
+    if hf_token:
+        print("✅ HF Token found and set")
+    else:
+        print("⚠️ No roberta_token found - model won't be pushed to Hub")
+    
     try:
         # Run training script
         process = subprocess.Popen(
-            [sys.executable, 'train.py'], 
-            stdout=subprocess.PIPE, 
+            [sys.executable, 'train.py'],
+            stdout=subprocess.PIPE,
             stderr=subprocess.STDOUT,
             universal_newlines=True,
             bufsize=1
@@ -32,65 +42,161 @@ def run_training(model_name):
                 break
         
         if process.returncode == 0:
-            output += f"\n\n✅ SUCCESS! Model saved to: https://huggingface.co/{model_name}"
+            if hf_token:
+                output += f"\n\n✅ SUCCESS! Model should be saved to: https://huggingface.co/{model_name}"
+            else:
+                output += f"\n\n✅ Training completed! Model saved locally (no HF token for Hub upload)"
         else:
             output += f"\n\n❌ Training failed with return code: {process.returncode}"
-            
+        
         return output
         
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
+def push_existing_model(model_name):
+    """Push already trained model to Hub"""
+    if not model_name.strip():
+        return "❌ Please enter a model name!"
+    
+    # Get token from environment
+    hf_token = os.environ.get('roberta_token')
+    
+    if not hf_token:
+        return "❌ roberta_token not found in environment!\nMake sure roberta_token is set in your Space secrets."
+    
+    try:
+        output = "🔐 Logging into Hugging Face Hub...\n"
+        login(token=hf_token)
+        output += "✅ Login successful!\n\n"
+        
+        output += "📂 Checking for trained model...\n"
+        # Check if model exists
+        if not os.path.exists("./model_output"):
+            return output + "❌ ./model_output directory not found!\nMake sure you've run training first."
+            
+        # Load your already-trained model
+        output += "📂 Loading trained model from ./model_output...\n"
+        model = AutoPeftModelForQuestionAnswering.from_pretrained("./model_output")
+        tokenizer = AutoTokenizer.from_pretrained("./model_output")
+        output += "✅ Model loaded successfully!\n\n"
+        
+        # Push to Hub
+        model_name = model_name.strip()
+        output += f"⬆️ Pushing model to Hub: {model_name}\n"
+        
+        model.push_to_hub(model_name, private=False)
+        tokenizer.push_to_hub(model_name, private=False)
+        
+        output += f"🎉 SUCCESS! Model pushed to: https://huggingface.co/{model_name}\n"
+        
+        # Also push training info if it exists
+        training_info_path = "./model_output/training_info.json"
+        if os.path.exists(training_info_path):
+            from huggingface_hub import upload_file
+            upload_file(
+                path_or_fileobj=training_info_path,
+                path_in_repo="training_info.json",
+                repo_id=model_name,
+                repo_type="model"
+            )
+            output += "📊 Training info also uploaded!\n"
+        
+        return output
+        
+    except Exception as e:
+        return f"❌ Error: {str(e)}\n\nCommon issues:\n- Invalid token\n- Model name already exists\n- Network issues"
+
 # Gradio Interface
 with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
     gr.Markdown("""
     # 🤖 RoBERTa CUAD Question Answering Trainer
     
-    This will train a RoBERTa model with LoRA on the CUAD dataset and save it to Hugging Face Hub.
-    
-    **Instructions:**
-    1. Enter your desired model name (format: `your-username/model-name`)
-    2. Click "Start Training"
-    3. Wait ~20-30 minutes for training to complete
-    4. Your model will be saved and publicly available!
+    Train a RoBERTa model with LoRA on CUAD dataset OR push an existing trained model to Hub.
     """)
     
-    with gr.Row():
-        with gr.Column():
-            model_name_input = gr.Textbox(
-                label="Model Name", 
-                placeholder="your-username/roberta-cuad-qa",
-                info="This will be your model's name on Hugging Face Hub"
-            )
+    with gr.Tabs():
+        with gr.TabItem("🚀 Train New Model"):
+            gr.Markdown("""
+            **Instructions:**
+            1. Enter your desired model name (format: `your-username/model-name`)
+            2. Click "Start Training"
+            3. Wait ~45-75 minutes for training to complete
+            """)
             
-            train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
+            with gr.Row():
+                with gr.Column():
+                    model_name_input = gr.Textbox(
+                        label="Model Name",
+                        placeholder="your-username/roberta-cuad-qa",
+                        info="This will be your model's name on Hugging Face Hub"
+                    )
+                    train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
+                    
+                with gr.Column():
+                    gr.Markdown("""
+                    **Training Details:**
+                    - Model: RoBERTa-base + LoRA
+                    - Dataset: CUAD (4000 samples → 20K after tokenization)
+                    - Time: ~45-75 minutes
+                    - GPU: T4 (free)
+                    """)
+            
+            train_output = gr.Textbox(
+                label="Training Output",
+                lines=25,
+                max_lines=50,
+                show_copy_button=True
+            )
             
-        with gr.Column():
+            train_btn.click(
+                fn=run_training,
+                inputs=model_name_input,
+                outputs=train_output,
+                show_progress=True
+            )
+        
+        with gr.TabItem("⬆️ Push Existing Model"):
             gr.Markdown("""
-            **Training Details:**
-            - Model: RoBERTa-base + LoRA
-            - Dataset: CUAD (2000 samples)
-            - Time: ~20-30 minutes
-            - GPU: T4 (free)
+            **Push Already Trained Model:**
+            If you've already trained a model and it's saved locally, use this to upload it to Hub.
             """)
-    
-    output_box = gr.Textbox(
-        label="Training Output", 
-        lines=25, 
-        max_lines=50,
-        show_copy_button=True
-    )
-    
-    train_btn.click(
-        fn=run_training, 
-        inputs=model_name_input, 
-        outputs=output_box,
-        show_progress=True
-    )
+            
+            with gr.Row():
+                with gr.Column():
+                    push_model_name = gr.Textbox(
+                        label="Model Name",
+                        placeholder="your-username/roberta-cuad-qa",
+                        info="Name for your model on Hugging Face Hub"
+                    )
+                    push_btn = gr.Button("⬆️ Push to Hub", variant="secondary", size="lg")
+                    
+                with gr.Column():
+                    gr.Markdown("""
+                    **Requirements:**
+                    - Model must be trained and saved in ./model_output/
+                    - roberta_token must be set in Space secrets
+                    - Takes ~30 seconds
+                    """)
+            
+            push_output = gr.Textbox(
+                label="Push Output",
+                lines=15,
+                show_copy_button=True
+            )
+            
+            push_btn.click(
+                fn=push_existing_model,
+                inputs=push_model_name,
+                outputs=push_output,
+                show_progress=True
+            )
     
     gr.Markdown("""
     ---
-    **Note:** Make sure your model name follows the format `username/model-name` and you have write permissions.
+    **Setup Required:**
+    - Set `roberta_token` in Space Settings → Repository secrets
+    - Get your token from: https://huggingface.co/settings/tokens (with Write permissions)
     """)
 
 if __name__ == "__main__":