tests/test_config.py

from korbyte.__main__ import _parser
from korbyte.config import (
    CORE_VOCAB_SIZE,
    DEFAULT_REPOSITORY_ID,
    SPECIAL_TOKEN_COUNT,
    TOTAL_VOCAB_SIZE,
)
from korbyte.special_tokens import special_token_strings


def test_vocabulary_budget_is_exact() -> None:
    assert CORE_VOCAB_SIZE == 128_000
    assert SPECIAL_TOKEN_COUNT == 256
    assert TOTAL_VOCAB_SIZE == 128_256


def test_special_tokens_are_stable_and_unique() -> None:
    tokens = special_token_strings()
    assert len(tokens) == SPECIAL_TOKEN_COUNT
    assert len(set(tokens)) == SPECIAL_TOKEN_COUNT
    assert tokens[:4] == (
        "<|begin_of_text|>",
        "<|end_of_text|>",
        "<|pad|>",
        "<|unk|>",
    )


def test_render_defaults_to_public_repository() -> None:
    args = _parser().parse_args(["render"])
    assert args.repo_id == DEFAULT_REPOSITORY_ID == "dawncr0w/KorByte-128K"