main.py

'''
Mode 1: Visual-Language Contrastive Pre-training (ImageBind-style)
- Freeze: encoder, decoder, class_head, bbox_head, count_head, caption_head
- Train: backbone + text encoder with 3-way contrastive loss:
  * view1 <-> view2: Visual self-agreement (masked pose views)
  * view1 <-> text: Cross-modal alignment
  * view2 <-> text: Cross-modal alignment
- Uses InfoNCE loss with in-batch negatives (temperature=0.07)
> python main.py --mode 1 --num_train_epochs 50 --output_dir checkpoints/mode1

Mode 2: Joint training (load mode 1 checkpoint, train everything)
- Unfreeze everything (backbone, encoder, decoder, all heads)
- Train all parameters with all losses (localization + captioning)
- Uses the contrastive pre-trained backbone and text embeddings
> python main.py --mode 2 --num_train_epochs 100 --mode1_checkpoint checkpoints/mode1/mode1_final --output_dir checkpoints/mode2
'''
import gc
import os
import torch
from typing import Optional
from dataclasses import dataclass, field

from transformers import (
    AutoTokenizer, DeformableDetrConfig,
    HfArgumentParser, TrainingArguments, 
    Trainer, EarlyStoppingCallback,
)
from loader import DVCDataset, trainer_collate_fn
from pdvc import DeformableDetrForObjectDetection
from captioners import LSTMCaptioner, MBartDecoderCaptioner
from config import *


def is_bfloat16_supported(): # Checks if the current device supports bfloat16
    return torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8


@dataclass
class ModelArguments:
    d_model: int = field(default=1024)
    encoder_layers: int = field(default=2)
    decoder_layers: int = field(default=2)
    encoder_attention_heads: int = field(default=8)
    decoder_attention_heads: int = field(default=8)
    encoder_n_points: int = field(default=4)
    decoder_n_points: int = field(default=4)
    num_feature_levels: int = field(default=4, metadata={'help': 'The number of input feature levels'})
    num_queries: int = field(default=30, metadata={'help': 'Maximum number of events a window can have'})
    num_labels: int = field(default=1, metadata={'help': 'Single foreground class for caption'})
    auxiliary_loss: bool = field(default=True, metadata={'help': 'The training step may spend a time in per-layer caption alignment and Hungarian matching'})
    # LOSS weights (Eq 3 / App C.2): (cls, giou, counter, caption) = (2, 4, 2, 2); L1 box loss disabled (bbox=0).
    class_cost: float = field(default=2, metadata={'help': 'LOSS weight of the classification (focal) term'})
    bbox_cost: float = field(default=0, metadata={'help': 'LOSS weight of the L1 box term (paper L_total has no L1 term -> 0)'})
    giou_cost: float = field(default=4, metadata={'help': 'LOSS weight of the GIoU term'})
    counter_cost: float = field(default=2, metadata={'help': 'LOSS weight of the event counter (BCE) term'})
    caption_cost: float = field(default=2, metadata={'help': 'LOSS weight of the captioning (NLL) term'})
    # HUNGARIAN MATCHING cost weights (App C.2 / Eq 13): (cls, L1, giou) = (1, 5, 2). These are DELIBERATELY
    # distinct from the loss weights above (matcher's job = stable assignment; loss's job = drive learning).
    # Previously the matcher reused the loss weights, so L1 (the paper's dominant matching term) was 0.
    match_class_cost: float = field(default=1, metadata={'help': 'MATCHER cost weight of the classification term'})
    match_bbox_cost: float = field(default=5, metadata={'help': 'MATCHER cost weight of the L1 box term (dominant)'})
    match_giou_cost: float = field(default=2, metadata={'help': 'MATCHER cost weight of the GIoU term'})
    focal_alpha: float = field(default=0.25)
    with_box_refine: bool = field(default=True, metadata={'help': 'Learnt (True) or Ground truth proposals (False, all losses except caption loss will be disabled)'})

    # Caption head parameters
    num_cap_layers: int = field(default=3)
    cap_dropout_rate: float = field(default=0.1)
    captioner_type: str = field(default='mbart', metadata={'help': 'Type of captioner to use (mbart or lstms)'})


@dataclass
class DataArguments:
    max_tries: int = field(default=20, metadata={'help': 'Maximum attempts to find a valid window with at least one event'})
    noise_rate: float = field(default=0.15, metadata={'help': 'Proportion of words to mask for noise injection during non-streaming training'})
    pose_augment: bool = field(default=False, metadata={'help': 'Apply pose augmentation during training'})
    stride_ratio: float = field(default=0.9, metadata={'help': 'Stride ratio for window sampling during validation/testing'})
    min_events: int = field(default=1, metadata={'help': 'Minimum number of events in a window'})
    max_events: int = field(default=10, metadata={'help': 'Maximum number of events in a window'})
    max_event_tokens: int = field(default=40, metadata={'help': 'Maximum number of tokens per event/caption'})
    max_window_tokens: int = field(default=128, metadata={'help': 'Maximum number of tokens in a window for non-streaming input'})
    load_by: str = field(default='window', metadata={'help': "Load data by 'window' or by 'video'"})


@dataclass
class CustomTrainingArguments(TrainingArguments):
    output_dir: str = field(default=CHECKPOINT_DIR, metadata={'help': 'Directory for checkpoints and logs'})
    num_train_epochs: float = field(default=100, metadata={'help': 'Total number of training epochs'})
    save_safetensors: bool = field(default=False, metadata={'help': 'Disable safe serialization to avoid the error'})
    
    # Data processing
    # auto_find_batch_size=True, # Find batch size that fit memory via exponential decay, avoiding CUDA OOM
    per_device_train_batch_size: int = field(default=32, metadata={'help': 'Effective batch size = per_device_train_batch_size x gradient_accumulation_steps x num_devices'})
    per_device_eval_batch_size: int = field(default=32, metadata={'help': 'Can be higher if greedy but should be smaller if using beam search'})
    dataloader_num_workers: int = field(default=4, metadata={'help': 'Number of subprocesses to use for data loading'})

    # Precision & optimization
    optim: str = field(default='adamw_torch_fused', metadata={'help': 'Choose optimizer'})
    weight_decay: float = field(default=1e-4, metadata={'help': 'Low since random windows already provide regularization'})
    fp16: bool = field(default=not is_bfloat16_supported(), metadata={'help': 'Use mixed precision training if supported'})
    bf16: bool = field(default=is_bfloat16_supported(), metadata={'help': 'Use bfloat16 (if supported) instead of fp16 for mixed precision training'})
    learning_rate: float = field(default=5e-4, metadata={'help': 'Linear decay learning rate'})
    # early_stopping_patience: int = field(default=10, metadata={'help': 'Early stopping patience by validation loss or Bleu'})
    ddp_find_unused_parameters: bool = field(default=False, metadata={'help': 'Avoid DDP overhead if all parameters are used'})
    max_grad_norm: float = field(default=1.0, metadata={'help': 'Gradient clipping to avoid exploding gradients'})
    
    # Reporting
    report_to: Optional[str] = field(default='none', metadata={'help': 'Whether to report to wandb/tensorboard/none'})
    logging_strategy: str = field(default='epoch')
    # eval_strategy: str = field(default='epoch', metadata={'help': 'Evaluate after each epoch'})
    
    # Saving
    save_strategy: str = field(default='epoch')
    save_total_limit: Optional[int] = field(default=1)
    # metric_for_best_model: Optional[str] = field(default='eval_loss', metadata={'help': 'Use validation loss/Bleu for early stopping'})
    # greater_is_better: Optional[bool] = field(default=False, metadata={'help': 'Lower loss / Higher Bleu is better'})
    # load_best_model_at_end: bool = field(default=True, metadata={'help': 'Load the best model based on validation loss/Bleu'})

    # Training mode control
    mode: int = field(default=2, metadata={'help': 'Training mode: 1 for contrastive pre-training, 2 for joint training'})
    mode1_checkpoint: Optional[str] = field(default=None, metadata={'help': 'Path to mode 1 checkpoint for mode 2 training'})


def freeze_module(module):
    for param in module.parameters():
        param.requires_grad = False

def unfreeze_module(module):
    for param in module.parameters():
        param.requires_grad = True
        
def print_trainable_parameters(model):
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f'Trainable parameters: {trainable/1e6:.2f}M / {total/1e6:.2f}M ({100*trainable/total:.1f}%)')
    
def handle_key_mismatches(state_dict, model_state):
    filtered_state = {}
    for k, v in state_dict.items():
        if k in model_state:
            if model_state[k].shape == v.shape: filtered_state[k] = v
            else: print(f'=> Skipping {k}: shape mismatch {v.shape} vs {model_state[k].shape}')
        else: print(f'=> Skipping {k}: not in model')
    return filtered_state
    

def main():
    parser = HfArgumentParser((ModelArguments, DataArguments, CustomTrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    if training_args.mode == 1 and model_args.with_box_refine == False:
        raise ValueError('with_box_refine must be True in mode 1 because GT boxes are needed for contrastive learning.')
    elif training_args.mode == 2 and training_args.mode1_checkpoint is None:
        print('Warning: Found no --mode1_checkpoint for mode 2. The model will be trained from scratch.')
        if model_args.with_box_refine == False:
            print('Warning: Directly use GT boxes for captioning, so all losses except caption loss will be disabled. Only use this for ablation studies.')
    
    # Data Loading
    tokenizer = AutoTokenizer.from_pretrained(TRIMMED_TOKENIZER_DIR)
    train_dataset = DVCDataset(
        split='train', tokenizer=tokenizer, max_tries=data_args.max_tries, noise_rate=data_args.noise_rate, pose_augment=data_args.pose_augment, 
        min_events=data_args.min_events, max_events=data_args.max_events, max_window_tokens=data_args.max_window_tokens, 
        max_event_tokens=data_args.max_event_tokens, load_by=data_args.load_by, seed=training_args.seed
    )
    # val_dataset = DVCDataset(
    #     split='val', tokenizer=tokenizer, pose_augment=False, stride_ratio=data_args.stride_ratio, 
    #     min_events=data_args.min_events, max_events=data_args.max_events, max_event_tokens=data_args.max_event_tokens, 
    #     max_window_tokens=data_args.max_window_tokens, load_by=data_args.load_by, seed=training_args.seed
    # )
    if getattr(training_args, 'local_rank', -1) in (-1, 0): # Only log sizes on the main process to avoid clutter in DDP
        print(f'\nTraining Mode: {training_args.mode}')
        print(f'Train dataset: {len(train_dataset)} samples')
        # print(f'Val dataset: {len(val_dataset)} samples')

    # Build weight dict based on mode
    if training_args.mode == 1: # Mode 1: Only contrastive loss (other losses zeroed)
        weight_dict = {
            'loss_ce': 0, 
            'loss_bbox': 0, 
            'loss_giou': 0, 
            'loss_counter': 0, 
            'loss_caption': 0,  # Only contrastive loss matters
        }
        contrastive_mode = True
    else: # Mode 2: All losses with balanced weights for joint training
        weight_dict = {
            'loss_ce': model_args.class_cost, 
            'loss_bbox': model_args.bbox_cost, 
            'loss_giou': model_args.giou_cost, 
            'loss_counter': model_args.counter_cost, 
            'loss_caption': model_args.caption_cost
        }
        contrastive_mode = False

    # Model Setup
    config = DeformableDetrConfig(
        d_model=model_args.d_model,
        encoder_layers=model_args.encoder_layers,
        decoder_layers=model_args.decoder_layers,
        encoder_attention_heads=model_args.encoder_attention_heads,
        decoder_attention_heads=model_args.decoder_attention_heads,
        encoder_n_points=model_args.encoder_n_points,
        decoder_n_points=model_args.decoder_n_points,
        activation_function='gelu',
        num_feature_levels=model_args.num_feature_levels,  # The number of input feature levels
        num_queries=model_args.num_queries,                # Maximum number of events a window can have
        num_labels=model_args.num_labels,                  # Single foreground class for caption
        auxiliary_loss=model_args.auxiliary_loss,          # The training step may spend a time in per-layer caption alignment and Hungarian matching
        # config.{class,bbox,giou}_cost feed the Hungarian MATCHER (pdvc.py / loss.py read them). Use the
        # paper's matching weights (1, 5, 2), NOT the loss weights — the loss weights live in weight_dict above.
        class_cost=model_args.match_class_cost,            # MATCHER classification cost weight
        bbox_cost=model_args.match_bbox_cost,              # MATCHER L1 cost weight (paper's dominant matching term)
        giou_cost=model_args.match_giou_cost,              # MATCHER GIoU cost weight
        focal_alpha=model_args.focal_alpha,
        with_box_refine=model_args.with_box_refine,        # Learnt (True) or Ground truth proposals (False)
    )
    model = DeformableDetrForObjectDetection(
        config=config,
        captioner_class=MBartDecoderCaptioner if model_args.captioner_type=='mbart' else LSTMCaptioner,
        vocab_size=tokenizer.vocab_size,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        decoder_start_token_id=tokenizer.lang_code_to_id[TGT_LANG],
        num_cap_layers=model_args.num_cap_layers,
        cap_dropout_rate=model_args.cap_dropout_rate,
        max_event_tokens=data_args.max_event_tokens,
        max_events=data_args.max_events,
        weight_dict=weight_dict,
        use_gt_boxes_for_caption=not model_args.with_box_refine, # No GT boxes needed in 2-stage training
        contrastive_mode=contrastive_mode,  # Enable contrastive learning in mode 1
    ) # IMPORTANT: Do not .to(device); Trainer handles device placement and DDP
    
    # Load mode 1 checkpoint for mode 2
    if training_args.mode == 2 and training_args.mode1_checkpoint is not None:
        checkpoint_path = os.path.join(training_args.mode1_checkpoint, 'pytorch_model.bin')
        if os.path.exists(checkpoint_path):
            print(f'Loading mode 1 (contrastive) checkpoint from: {checkpoint_path}')
            state_dict = torch.load(checkpoint_path, map_location='cpu')
            
            # Handle potential key mismatches due to contrastive_mode difference
            filtered_state = handle_key_mismatches(state_dict, model.state_dict())
            model.load_state_dict(filtered_state, strict=False)
            print(f'Loaded {len(filtered_state)}/{len(state_dict)} parameters from mode 1')
        else: raise FileNotFoundError(f'Mode 1 checkpoint not found: {checkpoint_path}')
    
    # Setup freezing based on mode
    print('\n' + '='*80)
    if training_args.mode == 1: # Contrastive pre-training, freeze everything except backbone + text encoder
        print('MODE 1: Contrastive Pre-training (backbone + text encoder trainable)')
        freeze_module(model) # Freeze everything first
        unfreeze_module(model.transformer.backbone)  # Train the backbone
        if model.text_embed is not None: unfreeze_module(model.text_embed) # Train text embedding
        if model.text_proj is not None: unfreeze_module(model.text_proj)   # Train text projection
    else: # Mode 2: Joint training - unfreeze everything
        print('MODE 2: Joint Training (all parameters trainable)')
        unfreeze_module(model) # Unfreeze everything
    print('='*80)
    
    if getattr(training_args, 'local_rank', -1) in (-1, 0): 
        print_trainable_parameters(model)
    
    # Move to device
    if training_args._n_gpu <= 1:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model.to(device)

    # Trainer Setup
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        # eval_dataset=val_dataset,
        data_collator=trainer_collate_fn,
        # callbacks=[EarlyStoppingCallback(early_stopping_patience=training_args.early_stopping_patience)],
    )
    trainer.train()
    
    # Save final model
    save_path = os.path.join(training_args.output_dir, f'mode{training_args.mode}_final')
    trainer.save_model(save_path)
    
    if getattr(training_args, 'local_rank', -1) in (-1, 0):
        print(f'\nMode {training_args.mode} training complete!')
        print(f'Model saved to: {save_path}')
        
        if training_args.mode == 1:
            print(f'\nTo continue with Mode 2 (joint training), run:')
            print(f'python main.py --mode 2 --mode1_checkpoint {save_path} --output_dir checkpoints/mode2')
        elif training_args.mode == 2:
            print(f'\nTraining complete! The model is ready for evaluation.')
    
    # Cleanup to free memory
    model.to('cpu')
    del tokenizer, train_dataset, model, training_args, trainer
    gc.collect()
    if torch.cuda.is_available():  torch.cuda.empty_cache()


if __name__ == '__main__':
    main()