Spaces:
Sleeping
Sleeping
| 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() |