app.py

#!/usr/bin/env python3
"""
Hugging Face Space App for Morphological Transformer Training
This creates a Gradio interface for training and using the models
"""

import gradio as gr
import os
import json
import subprocess
import time
from pathlib import Path
from typing import Dict, List, Optional

# Import your training modules
from scripts.hf_cloud_training import CloudTrainingConfig, CloudMorphologicalTrainer
from scripts.transformer import TagTransformer
from scripts.morphological_dataset import build_vocabulary, analyze_vocabulary

class TrainingInterface:
    """Gradio interface for training morphological transformers"""
    
    def __init__(self):
        self.config = CloudTrainingConfig()
        self.training_process = None
        self.training_logs = []
        
    def start_training(self, 
                      model_name: str,
                      dataset_name: str, 
                      run_number: str,
                      batch_size: int,
                      learning_rate: float,
                      max_epochs: int,
                      use_wandb: bool,
                      wandb_project: str) -> str:
        """Start training process"""
        
        if self.training_process and self.training_process.poll() is None:
            return "❌ Training is already in progress. Please wait for it to complete."
        
        # Update configuration
        self.config.model_name = model_name
        self.config.dataset_name = dataset_name
        self.config.run_number = run_number
        self.config.batch_size = batch_size
        self.config.learning_rate = learning_rate
        self.config.max_epochs = max_epochs
        self.config.wandb_project = wandb_project if use_wandb else None
        
        # Set environment variables
        env = os.environ.copy()
        env.update({
            'MODEL_NAME': model_name,
            'DATASET_NAME': dataset_name,
            'RUN_NUMBER': run_number,
            'WANDB_PROJECT': wandb_project if use_wandb else '',
            'HF_TOKEN': os.getenv('HF_TOKEN', ''),
        })
        
        try:
            # Start training process
            self.training_process = subprocess.Popen(
                ['python', 'scripts/hf_cloud_training.py'],
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
                universal_newlines=True
            )
            
            return f"🚀 Training started for {model_name} on {dataset_name} run {run_number}"
            
        except Exception as e:
            return f"❌ Error starting training: {str(e)}"
    
    def get_training_status(self) -> str:
        """Get current training status"""
        if not self.training_process:
            return "No training process running"
        
        if self.training_process.poll() is None:
            return "🔄 Training in progress..."
        else:
            return_code = self.training_process.returncode
            if return_code == 0:
                return "✅ Training completed successfully!"
            else:
                return f"❌ Training failed with return code {return_code}"
    
    def get_training_logs(self) -> str:
        """Get training logs"""
        if not self.training_process:
            return "No training process running"
        
        try:
            # Read available output
            output = self.training_process.stdout.read() if self.training_process.stdout else ""
            return output[-2000:] if output else "No logs available yet"
        except:
            return "Error reading logs"
    
    def stop_training(self) -> str:
        """Stop training process"""
        if self.training_process and self.training_process.poll() is None:
            self.training_process.terminate()
            return "🛑 Training stopped"
        else:
            return "No training process to stop"
    
    def list_available_models(self) -> str:
        """List available trained models"""
        models_dir = Path("/models")
        if not models_dir.exists():
            return "No models directory found"
        
        models = []
        for model_path in models_dir.iterdir():
            if model_path.is_dir():
                config_file = model_path / "config.json"
                if config_file.exists():
                    try:
                        with open(config_file) as f:
                            config = json.load(f)
                        models.append(f"📁 {model_path.name}: {config.get('model_type', 'unknown')}")
                    except:
                        models.append(f"📁 {model_path.name}: (config error)")
        
        if not models:
            return "No trained models found"
        
        return "\n".join(models)
    
    def test_model(self, model_name: str, input_text: str) -> str:
        """Test a trained model"""
        model_path = Path("/models") / model_name
        if not model_path.exists():
            return f"❌ Model {model_name} not found"
        
        try:
            # Load model configuration
            config_file = model_path / "config.json"
            with open(config_file) as f:
                config = json.load(f)
            
            # Load vocabularies
            src_vocab_file = model_path / "src_vocab.json"
            tgt_vocab_file = model_path / "tgt_vocab.json"
            
            if not (src_vocab_file.exists() and tgt_vocab_file.exists()):
                return "❌ Vocabulary files not found"
            
            with open(src_vocab_file) as f:
                src_vocab = json.load(f)
            with open(tgt_vocab_file) as f:
                tgt_vocab = json.load(f)
            
            # Create model
            model = TagTransformer(
                src_vocab_size=len(src_vocab),
                trg_vocab_size=len(tgt_vocab),
                embed_dim=config['embed_dim'],
                nb_heads=config['nb_heads'],
                src_hid_size=config['src_hid_size'],
                src_nb_layers=config['src_nb_layers'],
                trg_hid_size=config['trg_hid_size'],
                trg_nb_layers=config['trg_nb_layers'],
                dropout_p=config['dropout_p'],
                tie_trg_embed=config['tie_trg_embed'],
                label_smooth=config['label_smooth'],
                nb_attr=config.get('nb_attr', 0),
                src_c2i=src_vocab,
                trg_c2i=tgt_vocab,
                attr_c2i={},
            )
            
            # Load model weights
            import torch
            state_dict = torch.load(model_path / "pytorch_model.bin", map_location='cpu')
            model.load_state_dict(state_dict)
            model.eval()
            
            # Simple tokenization for testing
            input_tokens = input_text.strip().split()
            input_ids = [src_vocab.get(token, src_vocab['<UNK>']) for token in input_tokens]
            
            # Add BOS and EOS
            input_ids = [src_vocab['<BOS>']] + input_ids + [src_vocab['<EOS>']]
            
            # Convert to tensor
            input_tensor = torch.tensor([input_ids], dtype=torch.long)
            
            # Generate output (simplified)
            with torch.no_grad():
                # This is a simplified generation - you might want to implement proper inference
                output = model(input_tensor, None, input_tensor, None)
                output_ids = torch.argmax(output, dim=-1)
            
            # Convert back to text
            output_tokens = []
            for token_id in output_ids[0]:
                if token_id.item() in tgt_vocab.values():
                    token = [k for k, v in tgt_vocab.items() if v == token_id.item()][0]
                    if token not in ['<BOS>', '<EOS>', '<PAD>']:
                        output_tokens.append(token)
            
            return f"Input: {input_text}\nOutput: {' '.join(output_tokens)}"
            
        except Exception as e:
            return f"❌ Error testing model: {str(e)}"

# Create interface instance
interface = TrainingInterface()

# Create Gradio interface
with gr.Blocks(title="Morphological Transformer Training") as app:
    gr.Markdown("# 🚀 Morphological Transformer Training on Hugging Face")
    gr.Markdown("Train and test morphological reinflection models using TagTransformer architecture.")
    
    with gr.Tabs():
        # Training Tab
        with gr.Tab("🏋️ Training"):
            gr.Markdown("## Start Training")
            
            with gr.Row():
                with gr.Column():
                    model_name = gr.Textbox(
                        label="Model Name",
                        value="morphological-transformer",
                        placeholder="your-username/morphological-transformer"
                    )
                    dataset_name = gr.Dropdown(
                        label="Dataset",
                        choices=["10L_90NL", "50L_50NL", "90L_10NL"],
                        value="10L_90NL"
                    )
                    run_number = gr.Dropdown(
                        label="Run Number",
                        choices=["1", "2", "3"],
                        value="1"
                    )
                    
                with gr.Column():
                    batch_size = gr.Slider(
                        label="Batch Size",
                        minimum=8,
                        maximum=64,
                        value=32,
                        step=8
                    )
                    learning_rate = gr.Slider(
                        label="Learning Rate",
                        minimum=0.0001,
                        maximum=0.01,
                        value=0.001,
                        step=0.0001
                    )
                    max_epochs = gr.Slider(
                        label="Max Epochs",
                        minimum=10,
                        maximum=200,
                        value=100,
                        step=10
                    )
            
            with gr.Row():
                use_wandb = gr.Checkbox(label="Use Weights & Biases", value=True)
                wandb_project = gr.Textbox(
                    label="WandB Project",
                    value="morphological-transformer-cloud",
                    visible=True
                )
            
            with gr.Row():
                start_btn = gr.Button("🚀 Start Training", variant="primary")
                stop_btn = gr.Button("🛑 Stop Training", variant="secondary")
            
            training_status = gr.Textbox(
                label="Training Status",
                value="No training process running",
                interactive=False
            )
            
            training_logs = gr.Textbox(
                label="Training Logs",
                lines=10,
                interactive=False
            )
            
            # Event handlers
            start_btn.click(
                fn=interface.start_training,
                inputs=[model_name, dataset_name, run_number, batch_size, learning_rate, max_epochs, use_wandb, wandb_project],
                outputs=training_status
            )
            
            stop_btn.click(
                fn=interface.stop_training,
                outputs=training_status
            )
            
            # Auto-refresh status and logs
            app.load(
                fn=interface.get_training_status,
                outputs=training_status,
                every=5
            )
            
            app.load(
                fn=interface.get_training_logs,
                outputs=training_logs,
                every=10
            )
        
        # Models Tab
        with gr.Tab("📁 Models"):
            gr.Markdown("## Available Models")
            
            refresh_btn = gr.Button("🔄 Refresh Models")
            models_list = gr.Textbox(
                label="Trained Models",
                lines=10,
                interactive=False
            )
            
            refresh_btn.click(
                fn=interface.list_available_models,
                outputs=models_list
            )
            
            # Auto-refresh models list
            app.load(
                fn=interface.list_available_models,
                outputs=models_list
            )
        
        # Testing Tab
        with gr.Tab("🧪 Testing"):
            gr.Markdown("## Test Trained Models")
            
            with gr.Row():
                with gr.Column():
                    test_model_name = gr.Textbox(
                        label="Model Name",
                        placeholder="morphological-transformer_best"
                    )
                    test_input = gr.Textbox(
                        label="Input Text",
                        placeholder="Enter morphological features and word form",
                        lines=3
                    )
                    test_btn = gr.Button("🧪 Test Model", variant="primary")
                
                with gr.Column():
                    test_output = gr.Textbox(
                        label="Model Output",
                        lines=5,
                        interactive=False
                    )
            
            test_btn.click(
                fn=interface.test_model,
                inputs=[test_model_name, test_input],
                outputs=test_output
            )
        
        # Info Tab
        with gr.Tab("ℹ️ Info"):
            gr.Markdown("""
            ## About This Space
            
            This Space allows you to train morphological reinflection models using the TagTransformer architecture.
            
            ### Features:
            - 🏋️ **Training**: Train models on different datasets (10L_90NL, 50L_50NL, 90L_10NL)
            - 📁 **Model Management**: View and manage trained models
            - 🧪 **Testing**: Test trained models with custom inputs
            - 📊 **Monitoring**: Integration with Weights & Biases for experiment tracking
            
            ### Datasets:
            - **10L_90NL**: 10% labeled, 90% non-labeled data
            - **50L_50NL**: 50% labeled, 50% non-labeled data
            - **90L_10NL**: 90% labeled, 10% non-labeled data
            
            ### Usage:
            1. Go to the Training tab
            2. Configure your training parameters
            3. Start training
            4. Monitor progress in the logs
            5. Test your trained models in the Testing tab
            
            ### Requirements:
            - Data files must be mounted to `/data`
            - Hugging Face token for model upload (set as `HF_TOKEN` environment variable)
            - Weights & Biases token for monitoring (optional)
            """)

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