tests/conftest.py

"""Shared fixtures.

Tests are split into two tiers:

* **Fast tests** parse configuration, tokenizer files, and the chat template. They
  need no weights and no GPU, so they are safe for CI on every commit.
* **Heavy tests** load the 9.65 B checkpoint. They are marked ``slow`` and skipped
  unless ``PIKO_MODEL_PATH`` points at a local checkpoint or Hub id.

Set ``PIKO_CONFIG_PATH`` to a directory holding just the small json/jinja files to
run the fast tier against a checkout that has no weights.
"""

from __future__ import annotations

import json
import os
from pathlib import Path

import pytest

REPO_ROOT = Path(__file__).resolve().parents[1]
FIXTURES = REPO_ROOT / "tests" / "fixtures"


def pytest_configure(config: pytest.Config) -> None:
    config.addinivalue_line("markers", "slow: requires the full checkpoint and a GPU")


@pytest.fixture(scope="session")
def config_dir() -> Path:
    """Directory containing config.json and tokenizer files."""
    for candidate in (
        os.environ.get("PIKO_CONFIG_PATH"),
        os.environ.get("PIKO_MODEL_PATH"),
    ):
        if candidate and (Path(candidate) / "config.json").is_file():
            return Path(candidate)
    if (FIXTURES / "config.json").is_file():
        return FIXTURES
    pytest.skip("No config directory: set PIKO_CONFIG_PATH or PIKO_MODEL_PATH")


@pytest.fixture(scope="session")
def model_config(config_dir: Path) -> dict:
    return json.loads((config_dir / "config.json").read_text(encoding="utf-8"))


@pytest.fixture(scope="session")
def model_path() -> str:
    path = os.environ.get("PIKO_MODEL_PATH")
    if not path:
        pytest.skip("PIKO_MODEL_PATH is not set; skipping tests that load weights")
    return path


@pytest.fixture(scope="session")
def loaded_model(model_path: str):
    """Load the checkpoint once for the whole slow tier."""
    torch = pytest.importorskip("torch")
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required: CPU offload corrupts this architecture")
    pytest.importorskip("bitsandbytes")
    from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig

    model = AutoModelForMultimodalLM.from_pretrained(
        model_path,
        dtype=torch.bfloat16,
        device_map={"": 0},
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        ),
    )
    model.eval()
    processor = AutoProcessor.from_pretrained(model_path)
    return model, processor


@pytest.fixture(scope="session")
def receipt_image(tmp_path_factory: pytest.TempPathFactory) -> Path:
    """A deterministic rendered receipt, built without touching the network."""
    from evaluation.custom_suite.build_assets import receipt  # type: ignore

    path = tmp_path_factory.mktemp("assets") / "receipt.png"
    receipt(path)
    return path