src/mini_transformer/inference.py

"""Programmatic inference helper for mini_transformer."""

from __future__ import annotations

import warnings
from pathlib import Path

import torch
from hydra.utils import to_absolute_path
from omegaconf import DictConfig, OmegaConf
from transformers import PreTrainedTokenizerFast

from .configs import InferAppCfg, ModelCfg, TokenizerCfg
from .transformer import BasicEncoderDecoderTransformer
from .utils import check_tokenizer_model_compatibility, set_global_seed


def _resolve_device(device_str: str | None) -> torch.device:
    """Pick a device honoring configs while falling back to CPU if unavailable."""

    if not device_str:
        return torch.device("cpu")
    try:
        device = torch.device(device_str)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"Invalid runtime.device value: {device_str!r}") from exc

    if device.type == "cuda":
        if not torch.cuda.is_available():
            warnings.warn(
                "CUDA device requested but torch.cuda.is_available() is False; falling back to CPU.",
                RuntimeWarning,
                stacklevel=2,
            )
            return torch.device("cpu")
        if device.index is not None and device.index >= torch.cuda.device_count():
            warnings.warn(
                f"CUDA device index {device.index} out of range; using cuda:0 instead.",
                RuntimeWarning,
                stacklevel=2,
            )
            return torch.device("cuda:0")
    return device


_SESSION_CACHE: dict[
    tuple[str, ...], tuple[BasicEncoderDecoderTransformer, PreTrainedTokenizerFast]
] = {}


def _session_cache_key(
    model_cfg: ModelCfg,
    tokenizer_cfg: TokenizerCfg,
    tokenizer_path: Path,
    checkpoint_path: Path | None,
    device: torch.device,
) -> tuple[str, ...]:
    special_tokens = ",".join(tokenizer_cfg.special_tokens) if tokenizer_cfg.special_tokens else ""
    return (
        model_cfg.name,
        str(tokenizer_path),
        str(checkpoint_path) if checkpoint_path else "",
        str(device),
        str(model_cfg.vocab_size),
        str(model_cfg.d_model),
        str(model_cfg.num_heads),
        str(model_cfg.num_layers),
        str(model_cfg.d_ff),
        f"{model_cfg.dropout_rate}",
        str(model_cfg.max_seq_len),
        str(model_cfg.pad_id),
        str(model_cfg.bos_id),
        str(model_cfg.eos_id),
        tokenizer_cfg.name,
        str(tokenizer_cfg.vocab_size),
        str(tokenizer_cfg.max_seq_len),
        tokenizer_cfg.pad_token or "",
        tokenizer_cfg.bos_token or "",
        tokenizer_cfg.eos_token or "",
        tokenizer_cfg.unk_token or "",
        special_tokens,
    )


def _get_or_create_session(
    *,
    model_cfg: ModelCfg,
    tokenizer_cfg: TokenizerCfg,
    tokenizer_path: Path,
    checkpoint_path: Path | None,
    device: torch.device,
) -> tuple[BasicEncoderDecoderTransformer, PreTrainedTokenizerFast]:
    key = _session_cache_key(model_cfg, tokenizer_cfg, tokenizer_path, checkpoint_path, device)
    if key in _SESSION_CACHE:
        return _SESSION_CACHE[key]

    transformer = BasicEncoderDecoderTransformer(model_cfg)
    if checkpoint_path:
        ckpt = torch.load(checkpoint_path, map_location=device)
        state = ckpt.get("model_state_dict", ckpt)
        transformer.load_state_dict(state, strict=True)
    transformer.to(device)
    transformer.eval()

    tokenizer = PreTrainedTokenizerFast(
        tokenizer_file=str(tokenizer_path),
        bos_token=tokenizer_cfg.bos_token,
        eos_token=tokenizer_cfg.eos_token,
        unk_token=tokenizer_cfg.unk_token,
        pad_token=tokenizer_cfg.pad_token,
        model_max_length=tokenizer_cfg.max_seq_len,
    )

    _SESSION_CACHE[key] = (transformer, tokenizer)
    return transformer, tokenizer


def _preprocess_text(text: str) -> str:
    """Lowercase and ensure sentence ends with punctuation.

    Simple preprocessing used before tokenization during inference:
    - Lowercase the entire input.
    - If the last non-space character is not one of '.!?', append a period.
    """
    s = text.strip()
    if not s:
        return s
    s = s.lower()
    last = s[-1]
    if last not in ".!?":
        s += "."
    return s


def run_inference(cfg: DictConfig) -> list[str]:
    """Run inference using a composed Hydra configuration."""

    scfg_temp = OmegaConf.merge(OmegaConf.structured(InferAppCfg), cfg)
    scfg: InferAppCfg = OmegaConf.to_object(scfg_temp)
    model_cfg = ModelCfg(**vars(scfg.model))
    tokenizer_cfg = TokenizerCfg(**vars(scfg.tokenizer))

    set_global_seed(scfg.runtime.seed)
    check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg)

    device = _resolve_device(getattr(scfg.runtime, "device", None))

    if not tokenizer_cfg.path:
        raise FileNotFoundError(
            "Tokenizer path is empty. Update your config or set MINI_TRANSFORMER_TOKENIZER_PATH."
        )

    tokenizer_path = Path(to_absolute_path(tokenizer_cfg.path))
    if not tokenizer_path.is_file():
        raise FileNotFoundError(
            f"Tokenizer file not found: {tokenizer_path}. Check your model's config or environment variables."
        )

    checkpoint_path: Path | None = None
    if scfg.model.best_checkpoint_path:
        checkpoint_path = Path(to_absolute_path(scfg.model.best_checkpoint_path))

    transformer, tokenizer = _get_or_create_session(
        model_cfg=model_cfg,
        tokenizer_cfg=tokenizer_cfg,
        tokenizer_path=tokenizer_path,
        checkpoint_path=checkpoint_path,
        device=device,
    )

    text_input = _preprocess_text(scfg.input_text)
    if not text_input:
        raise SystemExit('Pass text like: input_text="hello world"')

    encoded = tokenizer(text_input, padding=True, truncation=True, return_tensors="pt")
    src_ids = encoded["input_ids"].to(device)
    src_padd_mask = (encoded["attention_mask"] == 0).to(device)

    tgt_ids = transformer.generate(
        src_ids,
        src_padd_mask,
        max_new_tokens=scfg.generation.max_new_tokens,
        temperature=scfg.generation.temperature,
        top_k=scfg.generation.top_k,
        top_p=scfg.generation.top_p,
        do_sample=scfg.generation.do_sample,
        presence_penalty=scfg.generation.presence_penalty,
        frequency_penalty=scfg.generation.frequency_penalty,
        no_repeat_ngram=scfg.generation.no_repeat_ngram,
        min_steps_before_eos=scfg.generation.min_steps_before_eos,
        seed=scfg.runtime.seed,
    )
    return tokenizer.batch_decode(tgt_ids.cpu(), skip_special_tokens=True)