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