loader.py

"""Model and tokenizer loading utilities."""

from __future__ import annotations

import gc
from dataclasses import dataclass
from typing import Any

import torch
from peft import AutoPeftModelForCausalLM, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.tokenization_utils_base import PreTrainedTokenizerBase

from config import HF_TOKEN, ModelSpec


@dataclass(slots=True)
class LoadedModelBundle:
    """Loaded model artifacts kept in memory."""

    spec: ModelSpec
    model: Any
    tokenizer: PreTrainedTokenizerBase
    source_kind: str


def _default_dtype() -> torch.dtype:
    if not torch.cuda.is_available():
        return torch.float32
    return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16


def _model_load_kwargs() -> dict[str, Any]:
    kwargs: dict[str, Any] = {
        "dtype": _default_dtype(),
        "low_cpu_mem_usage": True,
    }
    if HF_TOKEN:
        kwargs["token"] = HF_TOKEN
    if torch.cuda.is_available():
        kwargs["device_map"] = "auto"
    return kwargs


def _prepare_tokenizer(
    repo_id: str,
    *,
    fallback_repo_id: str | None = None,
    trust_remote_code: bool = False,
) -> PreTrainedTokenizerBase:
    errors: list[str] = []
    candidates = [repo_id]
    if fallback_repo_id and fallback_repo_id != repo_id:
        candidates.append(fallback_repo_id)

    for candidate in candidates:
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                candidate,
                token=HF_TOKEN,
                trust_remote_code=trust_remote_code,
            )
            if tokenizer.pad_token is None and tokenizer.eos_token is not None:
                tokenizer.pad_token = tokenizer.eos_token
            return tokenizer
        except Exception as exc:  # pragma: no cover - depends on remote repo state
            errors.append(f"{candidate}: {exc}")

    raise RuntimeError("Tokenizer load failed. " + " | ".join(errors))


def _get_peft_config(repo_id: str) -> PeftConfig | None:
    try:
        return PeftConfig.from_pretrained(repo_id, token=HF_TOKEN)
    except Exception:
        return None


def load_model_bundle(spec: ModelSpec) -> LoadedModelBundle:
    """Load a model and tokenizer for a single configured repository."""

    if not spec.repo_id:
        raise ValueError(
            f"{spec.name} is not configured. Set its repo_id in config.py or via "
            "the matching NL2BASH_* environment variable."
        )

    peft_config = _get_peft_config(spec.repo_id)
    if peft_config is not None:
        model = AutoPeftModelForCausalLM.from_pretrained(
            spec.repo_id,
            trust_remote_code=spec.trust_remote_code,
            **_model_load_kwargs(),
        )
        tokenizer = _prepare_tokenizer(
            spec.repo_id,
            fallback_repo_id=peft_config.base_model_name_or_path,
            trust_remote_code=spec.trust_remote_code,
        )
        source_kind = "adapter"
    else:
        model = AutoModelForCausalLM.from_pretrained(
            spec.repo_id,
            trust_remote_code=spec.trust_remote_code,
            **_model_load_kwargs(),
        )
        tokenizer = _prepare_tokenizer(
            spec.repo_id,
            trust_remote_code=spec.trust_remote_code,
        )
        source_kind = "full-model"

    model.eval()
    return LoadedModelBundle(
        spec=spec,
        model=model,
        tokenizer=tokenizer,
        source_kind=source_kind,
    )


def get_model_device(model: Any) -> torch.device:
    """Return the device that should receive tokenized inputs."""

    if hasattr(model, "device") and isinstance(model.device, torch.device):
        return model.device

    try:
        return next(model.parameters()).device
    except StopIteration:
        return torch.device("cpu")


def release_model_bundle(bundle: LoadedModelBundle) -> None:
    """Best-effort release of model resources after cache eviction."""

    del bundle.model
    del bundle.tokenizer
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()