src/utils.py

from __future__ import annotations

import json
import random
from pathlib import Path
from typing import Any

import numpy as np
import torch

from .config import Config
from .model import MultiMDMTransformer
from .tokenizer import CharacterTokenizer


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 resolve_device(requested: str) -> torch.device:
    if requested == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if requested == "cuda" and not torch.cuda.is_available():
        print("CUDA unavailable; falling back to CPU.")
        return torch.device("cpu")
    return torch.device(requested)


def build_model(config: Config, tokenizer: CharacterTokenizer) -> MultiMDMTransformer:
    return MultiMDMTransformer(
        vocab_size=tokenizer.vocab_size,
        clean_vocab_size=tokenizer.clean_vocab_size,
        seq_len=config.seq_len,
        num_masks=config.num_masks,
        d_model=config.d_model,
        nhead=config.nhead,
        num_layers=config.num_layers,
        dim_feedforward=config.dim_feedforward,
        dropout=config.dropout,
    )


def save_json(value: Any, path: str | Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, indent=2, ensure_ascii=False)


def save_checkpoint(
    path: str | Path,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    config: Config,
    tokenizer: CharacterTokenizer,
    step: int,
    metrics: dict[str, Any],
) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "config": config.to_dict(),
            "vocab": tokenizer.vocab,
            "num_masks": tokenizer.num_masks,
            "step": step,
            "metrics": metrics,
        },
        path,
    )


def load_checkpoint(
    path: str | Path, device: str | torch.device = "cpu"
) -> tuple[MultiMDMTransformer, CharacterTokenizer, Config, dict[str, Any]]:
    device = torch.device(device)
    checkpoint = torch.load(path, map_location=device, weights_only=False)
    config = Config.from_dict(checkpoint["config"])
    tokenizer = CharacterTokenizer(checkpoint["vocab"], checkpoint["num_masks"])
    model = build_model(config, tokenizer).to(device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    return model, tokenizer, config, checkpoint