Add comprehensive test suite — 72 passing tests covering all components

tests/test_tokenizer.py

@@ -0,0 +1,353 @@
+"""
+Comprehensive tests for domainTokenizer core library.
+72 tests covering: schemas, field tokenizers, predefined schemas,
+DomainTokenizerBuilder pipeline, and end-to-end HF encoding.
+
+Run: pytest tests/test_tokenizer.py -v
+"""
+
+import json
+import math
+import sys
+from datetime import datetime
+
+import numpy as np
+import pytest
+
+from domain_tokenizer.schema import DomainSchema, FieldSpec, FieldType, CALENDAR_FIELD_SIZES
+from domain_tokenizer.tokenizers.field_tokenizers import (
+    SignTokenizer, MagnitudeBucketTokenizer, DiscreteNumericalTokenizer,
+    CalendarTokenizer, CategoricalTokenizer, create_field_tokenizer,
+)
+from domain_tokenizer.tokenizers.domain_tokenizer import DomainTokenizerBuilder
+from domain_tokenizer.schemas.predefined import FINANCE_SCHEMA, ECOMMERCE_SCHEMA, HEALTHCARE_SCHEMA
+
+
+class TestFieldSpec:
+    def test_sign_field(self):
+        spec = FieldSpec("amount_sign", FieldType.SIGN)
+        assert spec.token_count == 2
+        assert spec.tokens_per_event == 1
+
+    def test_numerical_continuous_field(self):
+        spec = FieldSpec("amount", FieldType.NUMERICAL_CONTINUOUS, n_bins=21)
+        assert spec.token_count == 21
+
+    def test_numerical_discrete_field(self):
+        spec = FieldSpec("quantity", FieldType.NUMERICAL_DISCRETE, max_value=10)
+        assert spec.token_count == 12
+
+    def test_categorical_field(self):
+        spec = FieldSpec("event_type", FieldType.CATEGORICAL_FIXED, categories=["a", "b", "c"])
+        assert spec.token_count == 4
+
+    def test_temporal_field(self):
+        spec = FieldSpec("ts", FieldType.TEMPORAL, calendar_fields=["month", "dow", "dom", "hour"])
+        assert spec.token_count == 74
+
+    def test_text_field(self):
+        spec = FieldSpec("desc", FieldType.TEXT)
+        assert spec.token_count == 0
+
+    def test_custom_prefix(self):
+        spec = FieldSpec("amount", FieldType.NUMERICAL_CONTINUOUS, prefix="PRICE")
+        assert spec.prefix == "PRICE"
+
+    def test_categorical_requires_categories(self):
+        with pytest.raises(ValueError):
+            FieldSpec("event", FieldType.CATEGORICAL_FIXED)
+
+    def test_discrete_requires_max_value(self):
+        with pytest.raises(ValueError):
+            FieldSpec("qty", FieldType.NUMERICAL_DISCRETE)
+
+
+class TestDomainSchema:
+    def test_finance_token_count(self):
+        expected = 8 + 2 + 21 + 74
+        assert FINANCE_SCHEMA.special_token_count == expected
+
+    def test_finance_fixed_tokens(self):
+        assert FINANCE_SCHEMA.fixed_tokens_per_event == 7
+
+    def test_has_text_fields(self):
+        assert FINANCE_SCHEMA.has_text_fields is True
+
+    def test_text_field_names(self):
+        assert FINANCE_SCHEMA.text_field_names == ["description"]
+
+    def test_fittable_fields(self):
+        assert FINANCE_SCHEMA.fittable_field_names == ["amount"]
+
+    def test_get_field(self):
+        assert FINANCE_SCHEMA.get_field("amount").field_type == FieldType.NUMERICAL_CONTINUOUS
+
+    def test_get_field_missing(self):
+        assert FINANCE_SCHEMA.get_field("nonexistent") is None
+
+    def test_summary(self):
+        assert "finance" in FINANCE_SCHEMA.summary()
+
+
+class TestSignTokenizer:
+    def test_positive(self):
+        assert SignTokenizer("S")(79.99) == "[S_POS]"
+
+    def test_negative(self):
+        assert SignTokenizer("S")(-50.0) == "[S_NEG]"
+
+    def test_zero(self):
+        assert SignTokenizer("S")(0.0) == "[S_POS]"
+
+    def test_none(self):
+        assert SignTokenizer("S")(None) == "[S_POS]"
+
+    def test_nan(self):
+        assert SignTokenizer("S")(float("nan")) == "[S_POS]"
+
+    def test_vocab_size(self):
+        assert SignTokenizer("S").vocab_size == 2
+
+    def test_custom_labels(self):
+        tok = SignTokenizer("D", pos_label="CREDIT", neg_label="DEBIT")
+        assert tok(100) == "[D_CREDIT]"
+        assert tok(-100) == "[D_DEBIT]"
+
+
+class TestMagnitudeBucketTokenizer:
+    def setup_method(self):
+        self.tok = MagnitudeBucketTokenizer("A", n_bins=5)
+        self.tok.fit(np.array([1, 2, 5, 10, 20, 50, 100, 200, 500, 1000]))
+
+    def test_low(self):
+        assert self.tok(1.0) == "[A_00]"
+
+    def test_high(self):
+        assert self.tok(1000.0) == "[A_04]"
+
+    def test_negative_abs(self):
+        assert self.tok(50.0) == self.tok(-50.0)
+
+    def test_none(self):
+        assert self.tok(None) == "[A_00]"
+
+    def test_nan(self):
+        assert self.tok(float("nan")) == "[A_00]"
+
+    def test_vocab(self):
+        assert self.tok.vocab_size == 5
+
+    def test_not_fitted(self):
+        with pytest.raises(RuntimeError):
+            MagnitudeBucketTokenizer("X")(50.0)
+
+    def test_empty_fit(self):
+        with pytest.raises(ValueError):
+            MagnitudeBucketTokenizer("X").fit(np.array([]))
+
+    def test_nubank_21(self):
+        tok = MagnitudeBucketTokenizer("A", n_bins=21)
+        tok.fit(np.random.lognormal(3, 1, 10000))
+        assert tok.vocab_size == 21
+        for v in [0.01, 1.0, 100.0, 10000.0]:
+            assert tok(v) in tok.vocab
+
+    def test_serialization(self):
+        d = self.tok.to_dict()
+        tok2 = MagnitudeBucketTokenizer.from_dict(d)
+        assert tok2(50.0) == self.tok(50.0)
+
+
+class TestDiscreteNumericalTokenizer:
+    def test_normal(self):
+        assert DiscreteNumericalTokenizer("Q", max_value=10)(3) == "[Q_03]"
+
+    def test_zero(self):
+        assert DiscreteNumericalTokenizer("Q", max_value=10)(0) == "[Q_00]"
+
+    def test_max(self):
+        assert DiscreteNumericalTokenizer("Q", max_value=10)(10) == "[Q_10]"
+
+    def test_overflow(self):
+        assert DiscreteNumericalTokenizer("Q", max_value=10)(15) == "[Q_OVER]"
+
+    def test_negative(self):
+        assert DiscreteNumericalTokenizer("Q", max_value=10)(-5) == "[Q_00]"
+
+    def test_none(self):
+        assert DiscreteNumericalTokenizer("Q", max_value=10)(None) == "[Q_00]"
+
+    def test_vocab(self):
+        assert DiscreteNumericalTokenizer("Q", max_value=10).vocab_size == 12
+
+
+class TestCalendarTokenizer:
+    def test_full(self):
+        tok = CalendarTokenizer("T", fields=["month", "dow", "dom", "hour"])
+        tokens = tok(datetime(2025, 3, 15, 14, 30))
+        assert len(tokens) == 4
+        assert tokens[0] == "[T_MON_03]"
+        assert tokens[3] == "[T_HOUR_14]"
+
+    def test_string_input(self):
+        assert CalendarTokenizer("T", ["month"])("2025-03-15T14:30:00") == ["[T_MON_03]"]
+
+    def test_date_only(self):
+        tokens = CalendarTokenizer("T", ["month", "dow"])("2025-03-15")
+        assert tokens[0] == "[T_MON_03]"
+
+    def test_vocab_standard(self):
+        assert CalendarTokenizer("T", ["month", "dow", "dom", "hour"]).vocab_size == 74
+
+    def test_subset(self):
+        assert CalendarTokenizer("T", ["month", "dow"]).vocab_size == 19
+
+    def test_invalid(self):
+        with pytest.raises(ValueError):
+            CalendarTokenizer("T", ["invalid"])
+
+    def test_quarter(self):
+        tok = CalendarTokenizer("T", ["quarter"])
+        assert tok(datetime(2025, 1, 1)) == ["[T_Q1]"]
+        assert tok(datetime(2025, 10, 1)) == ["[T_Q4]"]
+
+
+class TestCategoricalTokenizer:
+    def test_known(self):
+        assert CategoricalTokenizer("E", ["view", "buy"])("buy") == "[E_001]"
+
+    def test_unknown(self):
+        assert CategoricalTokenizer("E", ["view", "buy"])("refund") == "[E_UNK]"
+
+    def test_none(self):
+        assert CategoricalTokenizer("E", ["view"])( None) == "[E_UNK]"
+
+    def test_vocab_unk(self):
+        tok = CategoricalTokenizer("E", ["a", "b"])
+        assert "[E_UNK]" in tok.vocab
+        assert tok.vocab_size == 3
+
+    def test_decode(self):
+        tok = CategoricalTokenizer("E", ["view", "buy"])
+        assert tok.decode_token("[E_000]") == "view"
+
+
+class TestFactory:
+    def test_sign(self):
+        assert isinstance(create_field_tokenizer(FieldSpec("s", FieldType.SIGN)), SignTokenizer)
+
+    def test_magnitude(self):
+        assert isinstance(create_field_tokenizer(FieldSpec("a", FieldType.NUMERICAL_CONTINUOUS)), MagnitudeBucketTokenizer)
+
+    def test_discrete(self):
+        assert isinstance(create_field_tokenizer(FieldSpec("q", FieldType.NUMERICAL_DISCRETE, max_value=10)), DiscreteNumericalTokenizer)
+
+    def test_calendar(self):
+        assert isinstance(create_field_tokenizer(FieldSpec("t", FieldType.TEMPORAL)), CalendarTokenizer)
+
+    def test_categorical(self):
+        assert isinstance(create_field_tokenizer(FieldSpec("c", FieldType.CATEGORICAL_FIXED, categories=["a"])), CategoricalTokenizer)
+
+    def test_text_none(self):
+        assert create_field_tokenizer(FieldSpec("d", FieldType.TEXT)) is None
+
+
+class TestPredefinedSchemas:
+    def test_finance(self):
+        assert FINANCE_SCHEMA.name == "finance"
+        assert len(FINANCE_SCHEMA.fields) == 4
+
+    def test_ecommerce(self):
+        assert ECOMMERCE_SCHEMA.name == "ecommerce"
+        assert len(ECOMMERCE_SCHEMA.fields) == 6
+
+    def test_healthcare(self):
+        assert HEALTHCARE_SCHEMA.name == "healthcare"
+        assert len(HEALTHCARE_SCHEMA.fields) == 6
+
+    def test_nubank_97(self):
+        domain_tokens = sum(f.token_count for f in FINANCE_SCHEMA.fields)
+        assert domain_tokens == 97
+
+
+class TestDomainTokenizerBuilder:
+    @pytest.fixture
+    def events(self):
+        return [
+            {"amount_sign": 79.99, "amount": 79.99,
+             "timestamp": datetime(2025, 3, 15, 14, 30), "description": "AMAZON"},
+            {"amount_sign": -200.0, "amount": -200.0,
+             "timestamp": datetime(2025, 3, 16, 9, 15), "description": "SALARY"},
+            {"amount_sign": 12.50, "amount": 12.50,
+             "timestamp": datetime(2025, 3, 17, 18, 45), "description": "UBER"},
+        ]
+
+    @pytest.fixture
+    def corpus(self):
+        return ["AMAZON", "SALARY", "UBER", "GROCERY", "NETFLIX"] * 20
+
+    def test_create(self):
+        assert not DomainTokenizerBuilder(FINANCE_SCHEMA).is_fitted
+
+    def test_fit(self, events):
+        b = DomainTokenizerBuilder(FINANCE_SCHEMA)
+        b.fit(events)
+        assert b.is_fitted
+
+    def test_tokenize_event(self, events):
+        b = DomainTokenizerBuilder(FINANCE_SCHEMA)
+        b.fit(events)
+        tokens = b.tokenize_event(events[0])
+        assert len(tokens) >= 7
+        assert tokens[0].startswith("[AMT_SIGN_")
+
+    def test_tokenize_sequence(self, events):
+        b = DomainTokenizerBuilder(FINANCE_SCHEMA)
+        b.fit(events)
+        tokens = b.tokenize_sequence(events)
+        assert tokens[0] == "[BOS]"
+        assert tokens[-1] == "[EOS]"
+        assert tokens.count(FINANCE_SCHEMA.event_separator) == 2
+
+    def test_build(self, events, corpus):
+        b = DomainTokenizerBuilder(FINANCE_SCHEMA)
+        b.fit(events)
+        hf = b.build(text_corpus=corpus, bpe_vocab_size=300)
+        assert hf.pad_token == "[PAD]"
+        assert hf.convert_tokens_to_ids("[AMT_SIGN_POS]") != hf.unk_token_id
+
+    def test_end_to_end(self, events, corpus):
+        b = DomainTokenizerBuilder(FINANCE_SCHEMA)
+        b.fit(events)
+        hf = b.build(text_corpus=corpus, bpe_vocab_size=300)
+        enc = b.encode_sequence(events, hf, max_length=128)
+        assert len(enc["input_ids"]) == 128
+        assert sum(1 for m in enc["attention_mask"] if m == 1) > 10
+
+    def test_stats(self, events):
+        b = DomainTokenizerBuilder(FINANCE_SCHEMA)
+        b.fit(events)
+        s = b.get_stats()
+        assert s["schema_name"] == "finance"
+        assert s["is_fitted"]
+
+    def test_unfitted_raises(self):
+        with pytest.raises(RuntimeError):
+            DomainTokenizerBuilder(FINANCE_SCHEMA).build()
+
+
+class TestEcommerceBuilder:
+    def test_full(self):
+        events = [
+            {"event_type": "view", "price": 29.99, "quantity": 1,
+             "category": "electronics", "timestamp": datetime(2025, 3, 15, 10, 0),
+             "product_title": "Mouse"},
+            {"event_type": "purchase", "price": 29.99, "quantity": 2,
+             "category": "electronics", "timestamp": datetime(2025, 3, 15, 10, 10),
+             "product_title": "Mouse"},
+        ]
+        b = DomainTokenizerBuilder(ECOMMERCE_SCHEMA)
+        b.fit(events)
+        hf = b.build(text_corpus=["Mouse", "Keyboard"] * 20, bpe_vocab_size=200)
+        enc = b.encode_sequence(events, hf, max_length=256)
+        assert sum(1 for m in enc["attention_mask"] if m == 1) > 10