app.py

"""
Gradio app for mmBERT training on HuggingFace Spaces.
Runs training in background and shows status via web interface.
"""

import gradio as gr
import subprocess
import threading
import time
from pathlib import Path

# Training status
training_status = {
    'stage': 'Starting...',
    'progress': 0,
    'logs': [],
    'complete': False,
    'error': None
}

def run_training():
    """Run training in background thread."""
    global training_status
    
    training_status['stage'] = 'Starting training...'
    training_status['logs'].append('Launching train.py...')
    
    try:
        # Run training script and capture output
        process = subprocess.Popen(
            ['python', 'train.py'],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1
        )
        
        # Read output line by line
        for line in iter(process.stdout.readline, ''):
            line = line.strip()
            if line:
                training_status['logs'].append(line)
                # Keep last 100 lines
                if len(training_status['logs']) > 100:
                    training_status['logs'] = training_status['logs'][-100:]
                
                # Update stage based on output
                if 'Loading data' in line:
                    training_status['stage'] = 'Loading data...'
                elif 'Loading model' in line:
                    training_status['stage'] = 'Loading model...'
                elif 'Starting training' in line:
                    training_status['stage'] = 'Training in progress...'
                elif 'Evaluating' in line:
                    training_status['stage'] = 'Evaluating...'
                elif 'Pushing to Hub' in line:
                    training_status['stage'] = 'Uploading model...'
                elif 'Done!' in line:
                    training_status['stage'] = 'Complete!'
                    training_status['complete'] = True
                elif 'loss' in line.lower():
                    training_status['stage'] = f'Training: {line[-80:]}'
        
        process.wait()
        
        if process.returncode == 0:
            training_status['stage'] = '✅ Training complete! Model uploaded.'
            training_status['complete'] = True
        else:
            training_status['stage'] = f'❌ Training failed with exit code {process.returncode}'
            training_status['error'] = f'Exit code: {process.returncode}'
            
    except Exception as e:
        training_status['stage'] = f'❌ Error: {str(e)}'
        training_status['error'] = str(e)
        training_status['logs'].append(f'ERROR: {str(e)}')

def get_status():
    """Get current training status."""
    logs = '\n'.join(training_status['logs'][-50:])
    status = f"## Status: {training_status['stage']}\n\n### Recent Logs:\n```\n{logs}\n```"
    return status

def refresh_status():
    """Refresh the status display."""
    return get_status()

# Start training in background thread
training_thread = threading.Thread(target=run_training, daemon=True)
training_thread.start()

# Create Gradio interface
with gr.Blocks(title="FTM Zone Trainer") as demo:
    gr.Markdown("# 🚀 FTM Email Zone Classifier Training")
    gr.Markdown("Training mmBERT model on 38,809 email examples...")
    
    status_display = gr.Markdown(value=get_status)
    
    refresh_btn = gr.Button("🔄 Refresh Status")
    refresh_btn.click(fn=refresh_status, outputs=status_display)

if __name__ == '__main__':
    demo.launch(server_name="0.0.0.0", server_port=7860)