src/train_nemo.py

"""
Fine-tuning script pentru gabrielpirlo/Sped_ParakeetRomanian_110M_TDT-CTC
folosind NVIDIA NeMo toolkit.

Model: EncDecHybridRNNTCTCBPEModel (FastConformer Hybrid TDT-CTC 110M)
Dataset: TTS-Romanian (datadriven-company/TTS-Romanian)
"""
import os
import argparse
from pathlib import Path

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from pytorch_lightning.loggers import TensorBoardLogger

import nemo.collections.asr as nemo_asr
from nemo.collections.asr.models import EncDecHybridRNNTCTCBPEModel
from nemo.core.config import hydra_runner
from nemo.utils import logging
from nemo.utils.exp_manager import exp_manager

import torch


def setup_datasets(model, train_manifest: str, val_manifest: str, batch_size: int = 16):
    """Configurează dataset-urile pentru antrenament și validare."""
    
    # Configurare train dataset
    model.cfg.train_ds.manifest_filepath = train_manifest
    model.cfg.train_ds.batch_size = batch_size
    model.cfg.train_ds.sample_rate = 16000
    model.cfg.train_ds.shuffle = True
    model.cfg.train_ds.num_workers = 8
    model.cfg.train_ds.pin_memory = True
    model.cfg.train_ds.max_duration = 20.0
    model.cfg.train_ds.min_duration = 0.1
    model.cfg.train_ds.bucketing_strategy = "synced_randomized"
    
    # Configurare validation dataset
    model.cfg.validation_ds.manifest_filepath = val_manifest
    model.cfg.validation_ds.batch_size = batch_size // 2
    model.cfg.validation_ds.sample_rate = 16000
    model.cfg.validation_ds.shuffle = False
    model.cfg.validation_ds.num_workers = 8
    model.cfg.validation_ds.pin_memory = True
    
    # Setup data loaders
    model.setup_training_data(model.cfg.train_ds)
    model.setup_validation_data(model.cfg.validation_ds)
    
    logging.info(f"Train dataset: {train_manifest}")
    logging.info(f"Validation dataset: {val_manifest}")


def setup_optimizer(model, lr: float = 2e-3, warmup_steps: int = 10000):
    """Configurează optimizer-ul și scheduler-ul."""
    
    # Optimizer AdamW
    model.cfg.optim.name = "adamw"
    model.cfg.optim.lr = lr
    model.cfg.optim.betas = [0.9, 0.98]
    model.cfg.optim.weight_decay = 1e-3
    
    # Scheduler Noam Annealing
    model.cfg.optim.sched.name = "NoamAnnealing"
    model.cfg.optim.sched.warmup_steps = warmup_steps
    model.cfg.optim.sched.min_lr = 1e-6
    model.cfg.optim.sched.d_model = 256  # embedding dim
    
    model.setup_optimization(optim_config=model.cfg.optim)


def train_model(
    model_path: str,
    train_manifest: str,
    val_manifest: str,
    output_dir: str,
    max_epochs: int = 10,
    batch_size: int = 16,
    lr: float = 2e-3,
    warmup_steps: int = 5000,
    precision: str = "bf16-mixed",
    accumulate_grad_batches: int = 4,
    checkpoint_dir: str = None,
):
    """Rulează antrenamentul."""
    
    # Verificare GPU
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA nu este disponibil! Antrenamentul necesită GPU.")
    
    device_name = torch.cuda.get_device_name(0)
    device_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
    logging.info(f"GPU: {device_name} ({device_memory:.1f} GB)")
    
    # Încărcare model
    logging.info(f"Încărcare model din: {model_path}")
    model = EncDecHybridRNNTCTCBPEModel.restore_from(model_path)
    
    # Configurare datasets
    setup_datasets(model, train_manifest, val_manifest, batch_size)
    
    # Configurare optimizer
    setup_optimizer(model, lr, warmup_steps)
    
    # Configurare PyTorch Lightning Trainer
    checkpoint_callback = ModelCheckpoint(
        dirpath=output_dir,
        filename="parakeet-romanian-{epoch:02d}-{val_wer:.2f}",
        monitor="val_wer",
        mode="min",
        save_top_k=3,
        save_last=True,
        verbose=True,
    )
    
    early_stop_callback = EarlyStopping(
        monitor="val_wer",
        min_delta=0.001,
        patience=5,
        mode="min",
        verbose=True,
    )
    
    tb_logger = TensorBoardLogger(save_dir=output_dir, name="lightning_logs")
    
    trainer = pl.Trainer(
        max_epochs=max_epochs,
        accelerator="gpu",
        devices=1,
        precision=precision,
        accumulate_grad_batches=accumulate_grad_batches,
        gradient_clip_val=1.0,
        enable_progress_bar=True,
        enable_model_summary=True,
        logger=tb_logger,
        callbacks=[checkpoint_callback, early_stop_callback],
        default_root_dir=output_dir,
    )
    
    # Resume din checkpoint dacă există
    ckpt_path = None
    if checkpoint_dir:
        last_ckpt = os.path.join(checkpoint_dir, "last.ckpt")
        if os.path.exists(last_ckpt):
            ckpt_path = last_ckpt
            logging.info(f"Resume din checkpoint: {ckpt_path}")
    
    # Antrenament
    logging.info("Începere antrenament...")
    trainer.fit(model, ckpt_path=ckpt_path)
    
    # Salvare model final
    final_model_path = os.path.join(output_dir, "parakeet-romanian-finetuned.nemo")
    model.save_to(final_model_path)
    logging.info(f"Model salvat: {final_model_path}")
    
    # Evaluare finală
    logging.info("Evaluare finală...")
    val_results = trainer.validate(model)
    if val_results:
        logging.info(f"Rezultate validare: {val_results}")
    
    return model, trainer


def main():
    parser = argparse.ArgumentParser(description="Fine-tuning Parakeet Romanian cu NeMo")
    parser.add_argument("--model_path", required=True, help="Cale către model .nemo")
    parser.add_argument("--train_manifest", required=True, help="Manifest train JSON")
    parser.add_argument("--val_manifest", required=True, help="Manifest validation JSON")
    parser.add_argument("--output_dir", default="/mnt/parakeet-training/outputs/nemo-finetune")
    parser.add_argument("--max_epochs", type=int, default=10)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--lr", type=float, default=2e-3)
    parser.add_argument("--warmup_steps", type=int, default=5000)
    parser.add_argument("--precision", default="bf16-mixed")
    parser.add_argument("--accumulate_grad_batches", type=int, default=4)
    parser.add_argument("--checkpoint_dir", default=None)
    args = parser.parse_args()
    
    os.makedirs(args.output_dir, exist_ok=True)
    
    train_model(
        model_path=args.model_path,
        train_manifest=args.train_manifest,
        val_manifest=args.val_manifest,
        output_dir=args.output_dir,
        max_epochs=args.max_epochs,
        batch_size=args.batch_size,
        lr=args.lr,
        warmup_steps=args.warmup_steps,
        precision=args.precision,
        accumulate_grad_batches=args.accumulate_grad_batches,
        checkpoint_dir=args.checkpoint_dir,
    )


if __name__ == "__main__":
    main()