tests/test_custom_classifier.py

"""Tests for custom HuggingFace classifier inference helpers."""

from __future__ import annotations

import json

from app.engines.semantic.custom_classifier import _prepare_model_inputs
from app.engines.semantic.custom_classifier import _resolve_tokenizer_source


class DistilBertLikeModel:
    """Model forward without token_type_ids, matching DistilBERT."""

    def forward(self, input_ids, attention_mask=None):
        return None


class BertLikeModel:
    """Model forward with token_type_ids support."""

    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
        return None


class KwargsModel:
    """Model forward that accepts arbitrary tokenizer fields."""

    def forward(self, **kwargs):
        return None


def test_prepare_model_inputs_strips_token_type_ids_for_distilbert():
    inputs = {
        "input_ids": "ids",
        "attention_mask": "mask",
        "token_type_ids": "segments",
    }

    prepared = _prepare_model_inputs(DistilBertLikeModel(), inputs)

    assert prepared == {
        "input_ids": "ids",
        "attention_mask": "mask",
    }
    assert "token_type_ids" in inputs


def test_prepare_model_inputs_keeps_supported_token_type_ids():
    inputs = {
        "input_ids": "ids",
        "attention_mask": "mask",
        "token_type_ids": "segments",
    }

    prepared = _prepare_model_inputs(BertLikeModel(), inputs)

    assert prepared == inputs


def test_prepare_model_inputs_keeps_token_type_ids_for_kwargs_model():
    inputs = {
        "input_ids": "ids",
        "attention_mask": "mask",
        "token_type_ids": "segments",
    }

    prepared = _prepare_model_inputs(KwargsModel(), inputs)

    assert prepared == inputs


def test_resolve_tokenizer_source_uses_local_tokenizer_json(tmp_path):
    model_dir = tmp_path / "custom_malicious_intent"
    model_dir.mkdir()
    (model_dir / "tokenizer.json").write_text("{}", encoding="utf-8")

    assert _resolve_tokenizer_source(model_dir) == str(model_dir)


def test_resolve_tokenizer_source_falls_back_to_training_base_model(tmp_path):
    model_dir = tmp_path / "custom_prompt_injection"
    model_dir.mkdir()
    (model_dir / "tokenizer_config.json").write_text("{}", encoding="utf-8")
    (model_dir / "training_metadata.json").write_text(
        json.dumps({"base_model": "distilbert-base-uncased"}),
        encoding="utf-8",
    )

    assert _resolve_tokenizer_source(model_dir) == "distilbert-base-uncased"