app.py

import gradio as gr
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 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,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            bufsize=1
        )
        
        output = ""
        while True:
            line = process.stdout.readline()
            if line:
                output += line
                print(line.strip())  # For real-time logs
            if process.poll() is not None:
                break
        
        if process.returncode == 0:
            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"
        # FIXED: Script 2 saves to ./cuad_lora_out, not ./model_output
        model_dir = "./cuad_lora_out"
        if not os.path.exists(model_dir):
            return output + f"❌ {model_dir} directory not found!\nMake sure you've run training first."
        
        # FIXED: Better error handling for model loading
        try:
            output += f"📂 Loading trained model from {model_dir}...\n"
            model = AutoPeftModelForQuestionAnswering.from_pretrained(model_dir)
            tokenizer = AutoTokenizer.from_pretrained(model_dir)
            output += "✅ Model loaded successfully!\n\n"
        except Exception as e:
            return output + f"❌ Failed to load model: {str(e)}\nMake sure the model was trained successfully."
        
        # Push to Hub
        model_name = model_name.strip()
        output += f"⬆️ Pushing model to Hub: {model_name}\n"
        
        try:
            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"
        except Exception as e:
            return output + f"❌ Failed to push model: {str(e)}\nCheck if model name already exists or token has write permissions."
        
        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\n- Token lacks write permissions"

# Gradio Interface
with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
    gr.Markdown("""
    # 🤖 RoBERTa CUAD Question Answering Trainer
    
    Train a RoBERTa model with LoRA on CUAD dataset OR push an existing trained model to Hub.
    
    **For AvocadoMuffin:** Your models will be saved as `AvocadoMuffin/your-model-name`
    """)
    
    with gr.Tabs():
        with gr.TabItem("🚀 Train New Model"):
            gr.Markdown("""
            **Instructions:**
            1. Enter your desired model name (format: `AvocadoMuffin/model-name`)
            2. Click "Start Training"
            3. Wait ~25-30 minutes for training to complete
            
            **Example model names:**
            - `AvocadoMuffin/roberta-cuad-qa`
            - `AvocadoMuffin/roberta-legal-qa-v1`
            """)
            
            with gr.Row():
                with gr.Column():
                    model_name_input = gr.Textbox(
                        label="Model Name",
                        placeholder="AvocadoMuffin/roberta-cuad-qa",
                        info="This will be your model's name on Hugging Face Hub",
                        value="AvocadoMuffin/"
                    )
                    train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
                    
                with gr.Column():
                    gr.Markdown("""
                    **Training Details:**
                    - Model: RoBERTa-base-squad2 + LoRA
                    - Dataset: CUAD (balanced, ~18K samples)
                    - Time: ~25-30 minutes on L4/T4
                    - GPU: T4 (free tier)
                    - Will auto-push to your HF profile
                    """)
            
            train_output = 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=train_output,
                show_progress=True
            )
        
        with gr.TabItem("⬆️ Push Existing Model"):
            gr.Markdown("""
            **Push Already Trained Model:**
            If you've already trained a model and it's saved locally, use this to upload it to Hub.
            """)
            
            with gr.Row():
                with gr.Column():
                    push_model_name = gr.Textbox(
                        label="Model Name",
                        placeholder="AvocadoMuffin/roberta-cuad-qa",
                        info="Name for your model on Hugging Face Hub",
                        value="AvocadoMuffin/"
                    )
                    push_btn = gr.Button("⬆️ Push to Hub", variant="secondary", size="lg")
                    
                with gr.Column():
                    gr.Markdown("""
                    **Requirements:**
                    - Model must be trained and saved in ./cuad_lora_out/
                    - 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("""
    ---
    **Setup Required for AvocadoMuffin:**
    - ✅ Set `roberta_token` in Space Settings → Repository secrets  
    - ✅ Get your token from: https://huggingface.co/settings/tokens (with Write permissions)
    - ✅ Your trained models will appear at: `https://huggingface.co/AvocadoMuffin/model-name`
    """)

if __name__ == "__main__":
    demo.launch()