tests/test_models.py

from types import SimpleNamespace
import unittest
from unittest.mock import patch

import torch

from utils import models


class FakeTokenizer:
    def __init__(self):
        self.calls = []

    def __call__(
        self,
        text,
        padding,
        truncation,
        max_length,
        return_tensors,
    ):
        self.calls.append(
            {
                'text': text,
                'padding': padding,
                'truncation': truncation,
                'max_length': max_length,
                'return_tensors': return_tensors,
            }
        )
        return {
            'input_ids': torch.tensor([[101, 102, 0]]),
            'attention_mask': torch.tensor([[1, 1, 0]]),
        }


class FakeModel:
    def __init__(self):
        self.device = torch.device('cpu')
        self.eval_called = False

    def to(self, device):
        self.device = torch.device(device)
        return self

    def eval(self):
        self.eval_called = True
        return self

    def __call__(self, **inputs):
        hidden_states = torch.tensor(
            [[[3.0, 0.0], [0.0, 4.0], [5.0, 12.0]]],
            device=self.device,
        )
        return SimpleNamespace(last_hidden_state=hidden_states)


class Qwen3EmbeddingTest(unittest.TestCase):
    def tearDown(self):
        models.get_embedding_model.cache_clear()

    def test_qwen3_model_is_available_without_changing_default(self):
        self.assertEqual(
            models.list_models[0],
            'sentence-transformers/paraphrase-multilingual-mpnet-base-v2',
        )
        self.assertIn(models.QWEN3_EMBEDDING_MODEL, models.list_models)

    def test_qwen3_uses_official_pooling_shape_and_normalization(self):
        tokenizer = FakeTokenizer()
        model = FakeModel()
        models.get_embedding_model.cache_clear()

        with (
            patch.object(models.AutoTokenizer, 'from_pretrained', return_value=tokenizer) as load_tokenizer,
            patch.object(models.AutoModel, 'from_pretrained', return_value=model) as load_model,
        ):
            embedding_model = models.get_embedding_model(models.QWEN3_EMBEDDING_MODEL)
            embedding = embedding_model('hello', pooling='cls')

        load_tokenizer.assert_called_once_with(
            models.QWEN3_EMBEDDING_MODEL,
            padding_side='left',
        )
        load_model.assert_called_once_with(models.QWEN3_EMBEDDING_MODEL)
        self.assertTrue(model.eval_called)
        self.assertEqual(tokenizer.calls[0]['max_length'], 8192)
        self.assertEqual(tuple(embedding.shape), (2,))
        torch.testing.assert_close(embedding.cpu(), torch.tensor([0.0, 1.0]))
        self.assertAlmostEqual(torch.linalg.vector_norm(embedding).item(), 1.0)


if __name__ == '__main__':
    unittest.main()