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()