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