train.py

import os
import dataclasses
import torch
import transformers
from transformers import Trainer, TrainingArguments, TrainerCallback
from peft import LoraConfig, get_peft_model, TaskType
from huggingface_hub import HfApi, login
import wandb
from dotenv import load_dotenv
from config import TrainConfig, ModelConfig
from model import MultiModalModel
from data import AudioTextDataset, DataCollator


class SamplePredictionCallback(TrainerCallback):
    """Every N steps, print ground-truth vs model-predicted transcript for a few samples."""

    def __init__(self, tokenizer, data_collator, train_dataset, sample_every_n_steps: int = 100, num_samples: int = 2, prompt: str = "Transcribe the following audio:"):
        self.tokenizer = tokenizer
        self.data_collator = data_collator
        self.train_dataset = train_dataset
        self.sample_every_n_steps = sample_every_n_steps
        self.num_samples = num_samples
        self.prompt = prompt
    def on_log(self, args, state, control, model=None, **kwargs):
        if state.global_step == 0 or state.global_step % self.sample_every_n_steps != 0:
            return
        if model is None:
            return
        model.eval()
        device = next(model.parameters()).device
        try:
            indices = [i % len(self.train_dataset) for i in range(self.num_samples)]
            batch = self.data_collator([self.train_dataset[i] for i in indices])
            audio_values = batch["audio_values"].to(device)
            labels_batch = batch["labels"]
            continuations = batch.get("continuation", [""] * audio_values.size(0))
            prompt_ids = self.tokenizer(self.prompt, return_tensors="pt", add_special_tokens=True).input_ids.to(device)
            prompt_ids = prompt_ids.expand(audio_values.size(0), -1)
            with torch.no_grad():
                gen_ids = model.generate(
                    input_ids=prompt_ids,
                    audio_values=audio_values,
                    max_new_tokens=120,
                    do_sample=False,
                    pad_token_id=self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
                )
            prompt_len = prompt_ids.size(1)
            
            # Create a wandb Table
            columns = ["Step", "Audio Index", "Ground Truth", "Prediction", "Continuation"]
            table = wandb.Table(columns=columns)
            
            print(f"\n[WandB] Logging sample predictions at step {state.global_step}")
            
            for i in range(audio_values.size(0)):
                gt_tokens = [t for t in labels_batch[i].tolist() if t != -100]
                gt_text = self.tokenizer.decode(gt_tokens, skip_special_tokens=True).strip()
                pred_text = self.tokenizer.decode(gen_ids[i][prompt_len:], skip_special_tokens=True).strip()
                
                cont_ref = continuations[i] if i < len(continuations) else ""
                
                # Add row to table
                table.add_data(state.global_step, i, gt_text, pred_text, cont_ref)
                
            # Log the table to wandb
            if wandb.run is not None:
                wandb.log({"sample_predictions": table}, step=state.global_step)
            else:
                print("Warning: WandB run not active, skipping logging.")

        except Exception as e:
            print(f"[SamplePredictionCallback] Error: {e}\n")
        finally:
            model.train()


import shutil
import glob
from transformers.trainer_utils import get_last_checkpoint

class AggressiveDeleteCallback(TrainerCallback):
    """
    Deletes ALL existing checkpoints in output_dir *before* saving a new one
    to ensure we don't run out of disk space. 
    Only keeps the one we are currently training on (in memory) effectively,
    but on disk we want 0 checkpoints just before save.
    
    WARNING: If save fails, we have NO checkpoints on disk. Risk accepted by user.
    """
    def __init__(self, output_dir):
        self.output_dir = output_dir

    def on_step_end(self, args, state, control, **kwargs):
        # Check if we are about to save
        # Trainer checks: if save_strategy == "steps" and global_step % save_steps == 0
        if args.save_strategy == "steps" and args.save_steps > 0:
            if state.global_step > 0 and state.global_step % args.save_steps == 0:
                # We are about to save. Delete old checkpoints.
                print(f"\n[AggressiveDeleteCallback] Step {state.global_step}: Deleting old checkpoints to free space before saving...")
                
                # Verify we aren't deleting something we just wrote (unlikely in on_step_end, save happens after)
                # But we might be resuming.
                
                ckpts = glob.glob(os.path.join(self.output_dir, "checkpoint-*"))
                for ckpt in ckpts:
                    try:
                        shutil.rmtree(ckpt)
                        print(f"  Deleted {ckpt}")
                    except Exception as e:
                        print(f"  Failed to delete {ckpt}: {e}")

def train():
    # Load environment variables
    load_dotenv()

    # Load Configs
    train_config = TrainConfig()
    model_config = ModelConfig()
    
    # Initialize WandB
    wandb.init(
        project=train_config.wandb_project,
        entity=train_config.wandb_entity,
        name=train_config.wandb_run_name,
        config=dataclasses.asdict(train_config),
    )

    
    # Initialize Tokenizer & Processor
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_config.text_model_id)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    processor = transformers.AutoProcessor.from_pretrained(model_config.audio_model_id)
    
    # Initialize Model
    model = MultiModalModel(model_config)
    
    # Apply LoRA if requested
    if train_config.use_lora:
        peft_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM, 
            inference_mode=False, 
            r=train_config.lora_r, 
            lora_alpha=train_config.lora_alpha, 
            lora_dropout=train_config.lora_dropout,
            target_modules=["q_proj", "v_proj"]
        )
        model.llm = get_peft_model(model.llm, peft_config)
        model.llm.print_trainable_parameters()
        
    # Dataset
    train_dataset = AudioTextDataset(train_config, processor, model_config, tokenizer)
    data_collator = DataCollator(processor, tokenizer)
    
    # Training Arguments (tuned for A100 80GB: bf16, larger batch, fast dataloader)
    training_args = TrainingArguments(
        output_dir=train_config.output_dir,
        per_device_train_batch_size=train_config.batch_size,
        gradient_accumulation_steps=train_config.accum_steps,
        learning_rate=train_config.learning_rate,
        lr_scheduler_type=train_config.lr_scheduler_type,
        num_train_epochs=train_config.num_epochs,
        max_steps=train_config.max_steps,
        bf16=train_config.use_bf16,
        gradient_checkpointing=train_config.gradient_checkpointing,
        dataloader_num_workers=train_config.dataloader_num_workers,
        dataloader_pin_memory=train_config.dataloader_pin_memory,
        logging_steps=train_config.log_steps,
        logging_first_step=True,
        logging_nan_inf_filter=True,
        save_steps=train_config.save_steps,
        save_total_limit=train_config.save_total_limit,
        eval_strategy="no",  # change if val set provided
        remove_unused_columns=False,  # Important because we have custom forward signature
        report_to="wandb",
        log_level="info",
        log_level_replica="info",
    )

    sample_callback = SamplePredictionCallback(
        tokenizer=tokenizer,
        data_collator=data_collator,
        train_dataset=train_dataset,
        sample_every_n_steps=train_config.sample_pred_every_steps,
        num_samples=2,
        prompt="Transcribe the following audio:",
    )
    
    aggressive_delete_callback = AggressiveDeleteCallback(train_config.output_dir)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        data_collator=data_collator,
        callbacks=[sample_callback, aggressive_delete_callback],
    )

    total_steps = train_config.max_steps
    print(f"\n>>> Training: max_steps={total_steps}, batch_size={train_config.batch_size}, "
          f"grad_accum={train_config.accum_steps} (effective batch={train_config.batch_size * train_config.accum_steps})")
    print(f">>> Sample predictions (GT vs predicted transcript) every {train_config.sample_pred_every_steps} steps.\n")

    # Resume from checkpoint if exists
    last_checkpoint = get_last_checkpoint(train_config.output_dir)
    if last_checkpoint is not None:
        print(f">>> Resuming from checkpoint: {last_checkpoint}")
        trainer.train(resume_from_checkpoint=last_checkpoint)
    else:
        trainer.train()
    
    # Save
    trainer.save_model(train_config.output_dir)
    tokenizer.save_pretrained(train_config.output_dir)
    processor.save_pretrained(train_config.output_dir)

    # Push to Hub
    if train_config.push_to_hub:
        print(f"\n>>> Pushing model to Hugging Face Hub: {train_config.hub_model_id}")
        if train_config.hub_token:
            login(token=train_config.hub_token)
        
        api = HfApi()
        
        # Create repo if needed
        # private=True by default for safety, user can adjust
        try:
            api.create_repo(repo_id=train_config.hub_model_id, private=train_config.hub_private_repo, exist_ok=True)
        except Exception as e:
            print(f"Warning: Could not create repo {train_config.hub_model_id}. Error: {e}")
        
        # Upload model folder
        try:
            api.upload_folder(
                folder_path=train_config.output_dir,
                repo_id=train_config.hub_model_id,
                repo_type="model",
            )
            
            # Upload code files to ensure custom model works
            for file in ["model.py", "config.py", "data.py", "inference.py"]:
                 if os.path.exists(file):
                      api.upload_file(
                           path_or_fileobj=file,
                           path_in_repo=file,
                           repo_id=train_config.hub_model_id,
                           repo_type="model",
                      )

            print(f">>> Successfully pushed to {train_config.hub_model_id}")
        except Exception as e:
            print(f"Error pushing to hub: {e}")

if __name__ == "__main__":
    train()