padoc/__init__.py

"""Public PaDoc API with heavy model imports loaded lazily."""

from __future__ import annotations

from importlib import import_module

from .constants import DEFAULT_FORK_TOKEN_MAP, DEFAULT_SPECIAL_TOKENS
from .data import DataPiece, Node, parse_data_piece, parse_node

__version__ = "0.1.0"

_LAZY = {
    "ForkTokenTrainer": ("padoc.modeling", "ForkTokenTrainer"),
    "SequentialPaDocEngine": ("padoc.transformers_infer", "SequentialPaDocEngine"),
    "preprocess_model": ("padoc.preprocess", "preprocess_model"),
}


def __getattr__(name: str):
    target = _LAZY.get(name)
    if target is None:
        raise AttributeError(f"module 'padoc' has no attribute {name!r}")
    module_name, attribute_name = target
    value = getattr(import_module(module_name), attribute_name)
    globals()[name] = value
    return value


__all__ = [
    "DEFAULT_FORK_TOKEN_MAP",
    "DEFAULT_SPECIAL_TOKENS",
    "DataPiece",
    "ForkTokenTrainer",
    "Node",
    "SequentialPaDocEngine",
    "parse_data_piece",
    "parse_node",
    "preprocess_model",
]