utils.py

"""Legacy convenience module that also exposes shared runtime helpers."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import torch
from torchvision import transforms
from transformers import AutoTokenizer

# Re-export model components for backward compatibility
from models import (
    BaseEncoder,
    BaseDecoder,
    EncoderCNN,
    EncoderViT,
    DecoderRNN,
    DecoderLSTM,
    DecoderTransformer,
    CaptioningModel,
)

__all__ = [
    "BaseEncoder",
    "BaseDecoder",
    "EncoderCNN",
    "EncoderViT",
    "DecoderRNN",
    "DecoderLSTM",
    "DecoderTransformer",
    "CaptioningModel",
    "build_default_transform",
    "resolve_device",
    "load_tokenizer",
    "load_checkpoint_bundle",
    "build_model_from_config",
    "load_model_from_checkpoint",
]


def build_default_transform(image_size: int = 224, augment: bool = True):
    """Standard training/eval image transforms.
    
    Args:
        image_size: Target image size
        augment: Whether to apply data augmentation (random crop/flip)
        
    Returns:
        Composed torchvision transforms
    """
    if augment:
        return transforms.Compose(
            [
                transforms.Resize(image_size + 32),
                transforms.RandomResizedCrop(image_size),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=(0.485, 0.456, 0.406),
                    std=(0.229, 0.224, 0.225),
                ),
            ]
        )
    return transforms.Compose(
        [
            transforms.Resize(image_size + 32),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=(0.485, 0.456, 0.406),
                std=(0.229, 0.224, 0.225),
            ),
        ]
    )


def resolve_device(device_arg: Optional[str] = None) -> torch.device:
    """Return torch.device, preferring CUDA when available."""
    if device_arg:
        return torch.device(device_arg)
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_tokenizer(tokenizer_name: str) -> AutoTokenizer:
    """Load a tokenizer and ensure pad token exists."""
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    if tokenizer.pad_token is None:
        if tokenizer.eos_token is not None:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({"pad_token": "[PAD]"})
    return tokenizer


def _infer_special_tokens(
    tokenizer: AutoTokenizer,
    checkpoint: Dict[str, Any],
) -> Tuple[int, int, int]:
    pad_token_id = checkpoint.get("pad_token_id", tokenizer.pad_token_id or 0)
    bos_token_id = checkpoint.get("bos_token_id")
    eos_token_id = checkpoint.get("eos_token_id")

    if bos_token_id is None:
        bos_token_id = tokenizer.cls_token_id or tokenizer.bos_token_id or pad_token_id
    if eos_token_id is None:
        eos_token_id = tokenizer.sep_token_id or tokenizer.eos_token_id or pad_token_id
    return pad_token_id, bos_token_id, eos_token_id


def load_checkpoint_bundle(
    checkpoint_path: str | Path,
    tokenizer_name: Optional[str] = None,
):
    """Load checkpoint metadata, config, and tokenizer in one place."""
    ckpt_path = Path(checkpoint_path).expanduser()
    if not ckpt_path.is_file():
        raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}")

    checkpoint = torch.load(ckpt_path, map_location="cpu")
    config = checkpoint.get("config")
    if config is None:
        raise ValueError(
            "Checkpoint is missing 'config'. Re-train with scripts/train_subset.py "
            "to embed model hyperparameters."
        )

    stored_name = checkpoint.get("tokenizer_name")
    effective_tokenizer = tokenizer_name or stored_name
    if effective_tokenizer is None:
        raise ValueError(
            "Tokenizer name not supplied via argument and missing from checkpoint metadata."
        )
    tokenizer = load_tokenizer(effective_tokenizer)
    pad_token_id, bos_token_id, eos_token_id = _infer_special_tokens(tokenizer, checkpoint)
    bundle = {
        "pad_token_id": pad_token_id,
        "bos_token_id": bos_token_id,
        "eos_token_id": eos_token_id,
        "tokenizer_name": effective_tokenizer,
        "vocab_size": checkpoint["vocab_size"],
    }
    return checkpoint, config, tokenizer, bundle


def build_model_from_config(
    config: Dict[str, Any],
    checkpoint: Dict[str, Any],
    *,
    pad_token_id: int,
    bos_token_id: int,
    eos_token_id: int,
    overrides: Optional[Dict[str, Any]] = None,
    strict: bool = True,
) -> CaptioningModel:
    """Instantiate CaptioningModel from checkpoint/config metadata."""
    model_kwargs: Dict[str, Any] = dict(
        embed_size=config["embed_size"],
        hidden_size=config["hidden_size"],
        vocab_size=checkpoint["vocab_size"],
        pad_token_id=pad_token_id,
        bos_token_id=bos_token_id,
        eos_token_id=eos_token_id,
        num_layers=config.get("num_layers", 1),
        dropout=config.get("dropout", 0.1),
        train_backbone=False,
        pretrained_backbone=config.get("pretrained_backbone", True),
        decoder_nonlinearity=config.get("decoder_nonlinearity", "tanh"),
        encoder_arch=config.get("encoder_arch", "cnn"),
        decoder_arch=config.get("decoder_arch", "rnn"),
        vit_model_name=config.get("vit_model_name", "google/vit-base-patch16-224-in21k"),
        vit_return_patches=config.get("vit_return_patches", False),
        transformer_num_layers=config.get("transformer_num_layers", 4),
        transformer_num_heads=config.get("transformer_num_heads", 8),
        transformer_ffn_size=config.get("transformer_ffn_size", 2048),
        transformer_dropout=config.get("transformer_dropout", 0.1),
        transformer_activation=config.get("transformer_activation", "gelu"),
        transformer_max_length=config.get("transformer_max_length", 64),
    )
    if overrides:
        model_kwargs.update(overrides)

    model = CaptioningModel(**model_kwargs)
    model.load_state_dict(checkpoint["model_state_dict"], strict=strict)
    return model


def load_model_from_checkpoint(
    checkpoint_path: str | Path,
    tokenizer_name: Optional[str] = None,
    *,
    model_overrides: Optional[Dict[str, Any]] = None,
    strict: bool = True,
) -> Tuple[CaptioningModel, AutoTokenizer, Dict[str, Any]]:
    """Load a CaptioningModel and tokenizer from disk."""
    checkpoint, config, tokenizer, bundle = load_checkpoint_bundle(
        checkpoint_path, tokenizer_name=tokenizer_name
    )
    model = build_model_from_config(
        config,
        checkpoint,
        pad_token_id=bundle["pad_token_id"],
        bos_token_id=bundle["bos_token_id"],
        eos_token_id=bundle["eos_token_id"],
        overrides=model_overrides,
        strict=strict,
    )
    metadata = {
        **bundle,
        "config": config,
        "checkpoint_path": str(Path(checkpoint_path).expanduser()),
    }
    return model, tokenizer, metadata