Spaces:
Sleeping
Sleeping
File size: 7,858 Bytes
c5a0569 947b7a8 c5a0569 947b7a8 c5a0569 947b7a8 c5a0569 947b7a8 c5a0569 947b7a8 c5a0569 947b7a8 c5a0569 947b7a8 3368d9b dbf32fb 3368d9b dbf32fb 947b7a8 dbf32fb 947b7a8 dbf32fb 947b7a8 c5a0569 947b7a8 dbf32fb c5a0569 947b7a8 dbf32fb 947b7a8 3368d9b dbf32fb 947b7a8 c5a0569 947b7a8 dbf32fb 947b7a8 3368d9b dbf32fb 947b7a8 c5a0569 947b7a8 c5a0569 947b7a8 c5a0569 947b7a8 dbf32fb 947b7a8 3368d9b 947b7a8 c5a0569 dbf32fb c5a0569 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 |
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() |