tests/test_model_loading.py

"""Configuration and loading contracts.

These assertions encode facts established in reports/repository_audit.md. If one
of them starts failing, either the checkpoint changed or the audit is stale —
both are worth knowing.
"""

from __future__ import annotations

import json
from pathlib import Path

import pytest

# --------------------------------------------------------------------------- #
# fast tier: no weights required
# --------------------------------------------------------------------------- #


def test_architecture_is_qwen3_5_conditional_generation(model_config: dict) -> None:
    assert model_config["architectures"] == ["Qwen3_5ForConditionalGeneration"]
    assert model_config["model_type"] == "qwen3_5"


def test_language_config_matches_audit(model_config: dict) -> None:
    text = model_config["text_config"]
    assert text["hidden_size"] == 4096
    assert text["num_hidden_layers"] == 32
    assert text["num_attention_heads"] == 16
    assert text["num_key_value_heads"] == 4
    assert text["head_dim"] == 256
    assert text["vocab_size"] == 248320
    assert text["max_position_embeddings"] == 262144


def test_layer_stack_is_hybrid(model_config: dict) -> None:
    """24 linear-attention layers, 8 full-attention, every 4th layer."""
    layers = model_config["text_config"]["layer_types"]
    assert len(layers) == 32
    assert layers.count("linear_attention") == 24
    assert layers.count("full_attention") == 8
    interval = model_config["text_config"]["full_attention_interval"]
    for index, kind in enumerate(layers):
        expected = "full_attention" if (index + 1) % interval == 0 else "linear_attention"
        assert kind == expected, f"layer {index} is {kind}, expected {expected}"


def test_vision_tower_projects_into_language_space(model_config: dict) -> None:
    vision = model_config["vision_config"]
    assert vision["depth"] == 27
    assert vision["hidden_size"] == 1152
    assert vision["patch_size"] == 16
    assert vision["out_hidden_size"] == model_config["text_config"]["hidden_size"]


def test_rope_is_native_not_scaled(model_config: dict) -> None:
    """The 262K context is architectural; no scaling factor should be present."""
    rope = model_config["text_config"]["rope_parameters"]
    assert rope["rope_type"] == "default"
    assert "factor" not in rope
    assert rope["mrope_section"] == [11, 11, 10]
    assert rope["mrope_interleaved"] is True


def test_no_remote_code_required(model_config: dict, config_dir: Path) -> None:
    assert "auto_map" not in model_config
    assert not list(config_dir.glob("*.py"))


def test_config_contains_no_absolute_paths(model_config: dict) -> None:
    """Regression guard for the absolute-path leak found in the original release."""
    blob = json.dumps(model_config)
    for marker in ("/mnt/", "/home/", "C:\\", "C:/", "\\Users\\"):
        assert marker not in blob, f"config.json leaks a local path containing {marker!r}"


def test_generation_config_defaults_to_greedy(config_dir: Path) -> None:
    path = config_dir / "generation_config.json"
    if not path.is_file():
        pytest.skip("generation_config.json not present")
    generation = json.loads(path.read_text(encoding="utf-8"))
    assert not generation.get("do_sample", False)
    assert "temperature" not in generation


def test_tokenizer_declares_vision_tokens(config_dir: Path) -> None:
    tokenizer = json.loads((config_dir / "tokenizer_config.json").read_text(encoding="utf-8"))
    special = tokenizer["model_specific_special_tokens"]
    assert special["image_token"] == "<|image_pad|>"
    assert special["vision_bos_token"] == "<|vision_start|>"
    assert tokenizer["padding_side"] == "left", "batched generation needs left padding"


def test_chat_template_present_and_handles_images(config_dir: Path) -> None:
    template = config_dir / "chat_template.jinja"
    assert template.is_file()
    body = template.read_text(encoding="utf-8")
    assert "<|vision_start|>" in body
    assert "<|image_pad|>" in body
    assert "add_generation_prompt" in body


# --------------------------------------------------------------------------- #
# slow tier: needs weights + GPU
# --------------------------------------------------------------------------- #


@pytest.mark.slow
def test_model_loads_without_trust_remote_code(loaded_model) -> None:
    model, _ = loaded_model
    assert type(model).__name__ == "Qwen3_5ForConditionalGeneration"


@pytest.mark.slow
def test_processor_loads(loaded_model) -> None:
    _, processor = loaded_model
    assert processor.__class__.__name__.startswith("Qwen3VL")


@pytest.mark.slow
def test_vision_tower_is_present_in_weights(loaded_model) -> None:
    model, _ = loaded_model
    vision_parameters = sum(p.numel() for p in model.model.visual.parameters())
    assert vision_parameters > 400_000_000