tests/test_decide_preprocessing.py

"""Bucket and marker checks for the real pinned GLiNER2.5-Decide tokenizer and schema."""
from pathlib import Path

import numpy as np
import pytest

from preprocessing import load_processor, prepare_decision
from runtime import decode

SOURCE = (
    Path.home()
    / ".cache/huggingface/hub/models--fastino--GLiNER2.5-Decide/snapshots/65624f1a0265b3f612bae66a2685a06b94a68a9d"
)
TASKS = {
    "intent": ["fyi", "request", "approval", "complaint", "newsletter", "security_alert"],
    "urgency": ["low", "normal", "high", "critical"],
    "route": ["support", "billing", "legal", "security", "finance", "archive"],
}
TEXT = "Please confirm the new retention rule is applied before Friday's audit."


@pytest.fixture(scope="module")
def processor():
    return load_processor(str(SOURCE))


def test_markers_point_at_label_tokens_per_head(processor):
    arrays = prepare_decision(processor, TEXT, TASKS, 128, 4, 8)
    assert arrays["marker_mask"][0].sum(axis=1).tolist() == [6, 4, 6, 0]
    label_id = processor.tokenizer.convert_tokens_to_ids("[L]")
    heads = arrays["marker_indices"][0]
    mask = arrays["marker_mask"][0] > 0.5
    assert np.all(arrays["input_ids"][0][heads[mask]] == label_id)


@pytest.mark.parametrize("bucket", [(32, 4, 8), (128, 2, 8), (128, 4, 5)])
def test_rejects_capacity_exceeded(processor, bucket):
    with pytest.raises(ValueError, match="bucket holds|decision heads"):
        prepare_decision(processor, TEXT, TASKS, *bucket)


def test_decode_matches_native_activation_rules():
    tasks = {
        "sentiment": ["positive", "negative"],
        "aspects": {"labels": ["battery", "screen", "price"], "multi_label": True, "cls_threshold": 0.4},
    }
    logits = np.array([[0.0, 2.0, -1e4], [3.0, -3.0, 0.0]], dtype=np.float32)
    result = decode(tasks, logits)
    assert result["sentiment"]["label"] == "negative"
    assert result["sentiment"]["confidence"] == pytest.approx(1 / (1 + np.exp(-2.0)))
    assert [entry["label"] for entry in result["aspects"]] == ["battery", "price"]


def test_decode_multi_label_falls_back_to_best_below_threshold():
    tasks = {"tags": {"labels": ["a", "b"], "multi_label": True, "cls_threshold": 0.99}}
    result = decode(tasks, np.array([[-1.0, 1.0]], dtype=np.float32))
    assert [entry["label"] for entry in result["tags"]] == ["b"]