Model/tests/conftest.py

"""
Shared pytest fixtures — a deliberately tiny MorphConfig so the full
test suite runs in seconds on CPU. Architecture ratios (GQA groups,
sliding window, YaRN scaling) are kept proportionally realistic; only
absolute sizes are shrunk.
"""
from __future__ import annotations

import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import torch

from config.model_config import MorphTextConfig, MorphConfig


@pytest.fixture
def tiny_text_config() -> MorphTextConfig:
    """~2M param LM — big enough to exercise every code path, small enough for CPU CI."""
    return MorphTextConfig(
        vocab_size=256,
        num_extra_tokens=16,
        hidden_size=64,
        intermediate_size=128,
        num_hidden_layers=4,
        num_attention_heads=8,
        num_key_value_heads=2,     # GQA 4:1
        head_dim=8,
        max_position_embeddings=512,
        rope_scaling_factor=2.0,
        qk_norm=True,
        use_sliding_window=True,
        sliding_window_size=16,    # small on purpose, to force the chunked fallback path
        torch_dtype="float32",     # CPU tests run in fp32
    )


@pytest.fixture
def tiny_config(tiny_text_config) -> MorphConfig:
    cfg = MorphConfig()
    cfg.text = tiny_text_config
    cfg.vision.projection_dim = tiny_text_config.hidden_size
    cfg.vision.hidden_size = 32
    cfg.vision.num_hidden_layers = 2
    cfg.vision.num_attention_heads = 4
    cfg.vision.intermediate_size = 64
    cfg.vision.image_size = 28
    cfg.vision.patch_size = 14        # 2x2 = 4 patches, +1 CLS = 5 tokens
    cfg.video.projection_dim = tiny_text_config.hidden_size
    return cfg


@pytest.fixture
def device() -> torch.device:
    return torch.device("cpu")