loaders.py

"""Helpers for loading the standalone EmbBERT bundle.

The simplified bundle ships one verified pretraining checkpoint together with
the tokenizer, configuration files, project library code, and runnable scripts.
The loader supports both the standalone ``EmbBERT/`` layout and the original
repository layout where the bundle lives in a subdirectory.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path


def bundle_root() -> Path:
    """Return the absolute path to the EmbBERT bundle directory."""

    return Path(__file__).resolve().parent


def repo_root() -> Path:
    """Return the directory that contains the runtime project files.

    Returns:
        The bundle root when it already contains ``lib/``. Otherwise, the
        parent directory is returned for the original repository layout.
    """

    for candidate in (bundle_root(), bundle_root().parent):
        if (candidate / "lib").exists():
            return candidate
    return bundle_root().parent


def _ensure_repo_on_path() -> None:
    """Add the repository root to ``sys.path`` if needed."""

    root = str(repo_root())
    if root not in sys.path:
        sys.path.insert(0, root)


def load_manifest() -> dict:
    """Load the bundle manifest.

    Returns:
        The decoded manifest dictionary.
    """

    manifest_path = bundle_root() / "manifest.json"
    return json.loads(manifest_path.read_text())


def build_tokenizer(max_length: int = 256):
    """Build a ``PreTrainedTokenizerFast`` from the bundled tokenizer file.

    Args:
        max_length: Maximum sequence length to expose through the tokenizer.

    Returns:
        A tokenizer configured with the special tokens used by this project.
    """

    _ensure_repo_on_path()

    from tokenizers import Tokenizer
    from transformers import PreTrainedTokenizerFast

    tokenizer_path = bundle_root() / "tokenizers" / "bpe_book_corpus_8192.json"
    tokenizer = Tokenizer.from_file(str(tokenizer_path))

    return PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        unk_token="[UNK]",
        sep_token="[SEP]",
        pad_token="[PAD]",
        cls_token="[CLS]",
        mask_token="[MASK]",
        max_length=max_length,
    )


def list_pretraining_checkpoints() -> list[str]:
    """Return bundled pretraining checkpoint names in ascending order."""

    manifest = load_manifest()
    return [
        artifact["name"]
        for artifact in manifest["artifacts"]
        if artifact["kind"] == "pretraining"
    ]


def load_pretraining_checkpoint(name: str = "checkpoint-616000"):
    """Load one bundled EmbBERT pretraining checkpoint.

    Args:
        name: Checkpoint directory name under ``checkpoints/pretraining``.

    Returns:
        A tuple ``(model, tokenizer)`` using the current repository classes.

    Raises:
        FileNotFoundError: If the requested checkpoint directory does not exist.
    """

    _ensure_repo_on_path()

    from lib.Models.EmbBERT import EmbBERT_Config
    from lib.Models.classifiers import PretrainingClassifier

    checkpoint_dir = bundle_root() / "checkpoints" / "pretraining" / name
    if not checkpoint_dir.exists():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_dir}")

    config = EmbBERT_Config.from_pretrained(str(checkpoint_dir))
    model = PretrainingClassifier.from_pretrained(str(checkpoint_dir), config=config)
    tokenizer = build_tokenizer(max_length=int(config.max_length))
    return model, tokenizer


def load_latest_pretraining_checkpoint():
    """Load the default bundled pretraining checkpoint.

    Returns:
        A tuple ``(model, tokenizer)`` for checkpoint ``616000``.
    """

    return load_pretraining_checkpoint("checkpoint-616000")


def load_legacy_sequence_metadata() -> dict:
    """Report that the simplified bundle no longer ships legacy artifacts.

    Raises:
        FileNotFoundError: Always raised because the simplified bundle only
            retains the ``checkpoint-616000`` pretraining artifact.
    """

    raise FileNotFoundError(
        "Legacy sequence metadata is not included in the simplified EmbBERT bundle."
    )