src/main.py

import json
import logging
import random
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional

import hydra
import matplotlib.pyplot as plt
import numpy as np
import torch
from omegaconf import OmegaConf
from sklearn.metrics import ConfusionMatrixDisplay, classification_report, confusion_matrix, f1_score
from torch import nn
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader

from dataset import MemeDataset
from models import (
    build_model,
    get_trainable_checkpoint,
    load_tokenizer_and_llama,
    load_trainable_checkpoint,
    resolve_dtype,
    trainable_parameters,
)
from train import collect_probabilities, optimize_threshold, run_epoch


logger = logging.getLogger(__name__)


def template_root() -> Path:
    return Path(__file__).resolve().parents[1]


def resolve_template_path(raw_path: str) -> Path:
    path = Path(str(raw_path))
    if path.is_absolute():
        return path
    return (template_root() / path).resolve()


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def make_loaders(cfg, tokenizer) -> tuple[DataLoader, DataLoader, DataLoader]:
    data_root = resolve_template_path(cfg.dataset.data_root)
    common_kwargs = {
        "data_root": data_root,
        "language": cfg.dataset.language,
        "tokenizer": tokenizer,
        "image_id_col": cfg.dataset.image_id_col,
        "text_col": cfg.dataset.text_col,
        "label_col": cfg.dataset.label_col,
        "max_length": cfg.training.max_length,
    }

    train_dataset = MemeDataset(
        split=cfg.dataset.train_split,
        oversample_positive_train=cfg.training.oversample_positive_train,
        **common_kwargs,
    )
    dev_dataset = MemeDataset(
        split=cfg.dataset.dev_split,
        oversample_positive_train=0,
        **common_kwargs,
    )
    test_dataset = MemeDataset(
        split=cfg.dataset.test_split,
        oversample_positive_train=0,
        **common_kwargs,
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=cfg.training.batch_size,
        shuffle=True,
        drop_last=True,
        num_workers=cfg.training.num_workers,
    )
    dev_loader = DataLoader(
        dev_dataset,
        batch_size=cfg.training.batch_size,
        shuffle=False,
        drop_last=False,
        num_workers=cfg.training.num_workers,
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=cfg.training.batch_size,
        shuffle=False,
        drop_last=False,
        num_workers=cfg.training.num_workers,
    )
    return train_loader, dev_loader, test_loader


def make_output_dirs(cfg) -> Dict[str, Path]:
    output_root = resolve_template_path(cfg.dataset.output_root)
    paths = {
        "predictions": output_root / "predictions" / cfg.dataset.language / "fusion",
        "metrics": output_root / "predictions" / cfg.dataset.language / "fusion",
        "models": output_root / "trained_model" / cfg.dataset.language / "fusion",
    }
    for path in paths.values():
        path.mkdir(parents=True, exist_ok=True)
    return paths


def save_run_artifacts(
    cfg,
    paths: Dict[str, Path],
    model,
    fusion: str,
    initial_lr: float,
    final_lr: float,
    best_epoch: int,
    best_threshold: float,
    best_dev_f1: float,
    threshold_scores: Dict[float, float],
    test_probs: np.ndarray,
    test_labels: np.ndarray,
    test_ids: np.ndarray,
    test_loss: float,
    test_macro_f1: float,
    test_accuracy: float,
) -> None:
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_name = (
        f"fusemd_{cfg.dataset.language}_{fusion}"
        f"_lr{initial_lr}_epoch{best_epoch}_bs{cfg.training.batch_size}_{timestamp}"
    )

    predictions = (test_probs >= best_threshold).astype(float)
    predictions_path = paths["predictions"] / f"{run_name}.csv"
    metrics_json_path = paths["metrics"] / f"{run_name}.json"
    metrics_txt_path = paths["metrics"] / f"{run_name}.txt"
    checkpoint_path = paths["models"] / f"{run_name}.pth"

    np.savetxt(
        predictions_path,
        np.column_stack(
            [
                test_ids.astype(int),
                test_probs,
                predictions.astype(int),
                test_labels.astype(int),
            ]
        ),
        delimiter=",",
        header="image_id,probability,prediction,true_label",
        comments="",
    )

    checkpoint = {
        **get_trainable_checkpoint(model),
        "fusion_method": fusion,
        "learning_rate": initial_lr,
        "final_learning_rate": final_lr,
        "best_epoch": best_epoch,
        "best_threshold": best_threshold,
        "best_dev_macro_f1": best_dev_f1,
        "threshold_scores": threshold_scores,
        "test_macro_f1": test_macro_f1,
        "test_accuracy": test_accuracy,
        "test_loss": test_loss,
        "config": OmegaConf.to_container(cfg, resolve=True),
    }
    torch.save(checkpoint, checkpoint_path)
    stable_checkpoint_path = resolve_template_path(cfg.dataset.finetuned_model_path)
    stable_checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(checkpoint, stable_checkpoint_path)

    report = classification_report(test_labels, predictions, digits=5, zero_division=0)
    matrix = confusion_matrix(test_labels, predictions)
    metrics = {
        "run_name": run_name,
        "date_time": datetime.now().isoformat(),
        "fusion_method": fusion,
        "initial_learning_rate": initial_lr,
        "final_learning_rate": final_lr,
        "best_epoch": best_epoch,
        "best_threshold": best_threshold,
        "best_dev_macro_f1": best_dev_f1,
        "threshold_scores": threshold_scores,
        "test_macro_f1": test_macro_f1,
        "test_accuracy": test_accuracy,
        "test_loss": test_loss,
        "confusion_matrix": matrix.tolist(),
        "classification_report": report,
        "config": OmegaConf.to_container(cfg, resolve=True),
    }

    with open(metrics_json_path, "w", encoding="utf-8") as file:
        json.dump(metrics, file, ensure_ascii=False, indent=2)

    with open(metrics_txt_path, "w", encoding="utf-8") as file:
        file.write(f"Run: {run_name}\n")
        file.write(f"Date and time: {metrics['date_time']}\n")
        file.write(f"Fusion: {fusion}\n")
        file.write(f"Initial LR: {initial_lr}\n")
        file.write(f"Final LR: {final_lr}\n")
        file.write(f"Best epoch: {best_epoch}\n")
        file.write(f"Best threshold: {best_threshold}\n")
        file.write(f"Dev macro-F1: {best_dev_f1:.5f}\n")
        file.write(f"Test macro-F1: {test_macro_f1:.5f}\n")
        file.write(f"Test accuracy: {test_accuracy:.5f}\n")
        file.write(f"Test loss: {test_loss:.5f}\n\n")
        file.write(str(matrix))
        file.write("\n\n")
        file.write(report)

    if cfg.training.save_confusion_matrix_png:
        display = ConfusionMatrixDisplay(matrix)
        display.plot()
        plt.title(f"{cfg.dataset.language} {fusion} confusion matrix")
        plt.savefig(paths["metrics"] / f"{run_name}_confusion_matrix.png", bbox_inches="tight", dpi=200)
        plt.close()

    logger.info("Saved checkpoint: %s", checkpoint_path)
    logger.info("Updated default inference checkpoint: %s", stable_checkpoint_path)
    logger.info("Saved predictions: %s", predictions_path)
    logger.info("Saved metrics: %s", metrics_json_path)


def train_one_run(
    cfg,
    tokenizer,
    llama,
    train_loader: DataLoader,
    dev_loader: DataLoader,
    test_loader: DataLoader,
    device: torch.device,
    fusion: str,
    learning_rate: float,
    paths: Dict[str, Path],
) -> None:
    logger.info("Starting run with fusion=%s learning_rate=%s", fusion, learning_rate)
    model = build_model(llama, cfg.model.image_model, fusion, cfg.model.embedding_dim).to(device)

    optimizer = torch.optim.Adam(trainable_parameters(model), lr=learning_rate, eps=cfg.training.optimizer_eps)
    scheduler = ReduceLROnPlateau(
        optimizer,
        factor=cfg.training.scheduler_factor,
        patience=cfg.training.scheduler_patience,
        threshold=cfg.training.scheduler_threshold,
    )
    criterion = nn.BCEWithLogitsLoss()

    best_dev_loss = float("inf")
    best_epoch = 0
    final_lr = learning_rate
    epochs_without_improvement = 0
    best_state: Optional[Dict[str, object]] = None

    for epoch in range(1, cfg.training.max_epochs + 1):
        train_loss = run_epoch(
            model,
            train_loader,
            criterion,
            device,
            optimizer=optimizer,
            description=f"Train epoch {epoch}",
        )
        with torch.no_grad():
            dev_loss = run_epoch(
                model,
                dev_loader,
                criterion,
                device,
                optimizer=None,
                description=f"Dev epoch {epoch}",
            )

        scheduler.step(dev_loss)
        current_lr = optimizer.param_groups[0]["lr"]
        logger.info(
            "Epoch %s train_loss=%.5f dev_loss=%.5f lr=%.8f",
            epoch,
            train_loss,
            dev_loss,
            current_lr,
        )

        if dev_loss < best_dev_loss:
            best_dev_loss = dev_loss
            best_epoch = epoch
            final_lr = current_lr
            epochs_without_improvement = 0
            best_state = get_trainable_checkpoint(model)
        else:
            epochs_without_improvement += 1

        if epochs_without_improvement >= cfg.training.early_stopping_patience:
            logger.info("Early stopping at epoch %s", epoch)
            break

    if best_state is not None:
        load_trainable_checkpoint(model, best_state)
        model.to(device)

    dev_probs, dev_labels, _ = collect_probabilities(model, dev_loader, device)
    best_threshold, best_dev_f1, threshold_scores = optimize_threshold(
        dev_probs,
        dev_labels,
        cfg.training.threshold_start,
        cfg.training.threshold_stop,
        cfg.training.threshold_step,
    )

    test_probs, test_labels, test_ids = collect_probabilities(model, test_loader, device)
    test_preds = (test_probs >= best_threshold).astype(float)
    test_macro_f1 = f1_score(test_labels, test_preds, average="macro", zero_division=0)
    test_accuracy = float((test_preds == test_labels).mean())

    with torch.no_grad():
        test_loss = run_epoch(
            model,
            test_loader,
            criterion,
            device,
            optimizer=None,
            description="Test loss",
        )

    save_run_artifacts(
        cfg=cfg,
        paths=paths,
        model=model,
        fusion=fusion,
        initial_lr=learning_rate,
        final_lr=final_lr,
        best_epoch=best_epoch,
        best_threshold=best_threshold,
        best_dev_f1=best_dev_f1,
        threshold_scores=threshold_scores,
        test_probs=test_probs,
        test_labels=test_labels,
        test_ids=test_ids,
        test_loss=test_loss,
        test_macro_f1=test_macro_f1,
        test_accuracy=test_accuracy,
    )

    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


@hydra.main(version_base=None, config_path="./configs", config_name="config")
def main(cfg) -> None:
    logger.info(OmegaConf.to_yaml(cfg, resolve=True))
    set_seed(cfg.training.seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = resolve_dtype(cfg.model.dtype)
    use_8bit = bool(cfg.model.use_8bit and device.type == "cuda")
    if cfg.model.use_8bit and device.type != "cuda":
        logger.warning("8-bit quantization requested, but CUDA is unavailable. Continuing without 8-bit.")

    tokenizer, llama = load_tokenizer_and_llama(cfg.model.text_model, use_8bit, dtype, device)
    train_loader, dev_loader, test_loader = make_loaders(cfg, tokenizer)
    paths = make_output_dirs(cfg)

    for fusion in cfg.model.fusion_methods:
        for learning_rate in cfg.training.learning_rates:
            train_one_run(
                cfg=cfg,
                tokenizer=tokenizer,
                llama=llama,
                train_loader=train_loader,
                dev_loader=dev_loader,
                test_loader=test_loader,
                device=device,
                fusion=fusion,
                learning_rate=float(learning_rate),
                paths=paths,
            )


if __name__ == "__main__":
    main()