Add field_tokenizers.py — Sign, MagnitudeBucket, Calendar, Categorical, DiscreteNumerical tokenizers

src/domain_tokenizer/tokenizers/field_tokenizers.py

@@ -0,0 +1,335 @@
+"""
+Per-field tokenizers for domain-specific data.
+
+Each tokenizer converts a single field value into one or more token strings.
+These are the building blocks assembled by DomainTokenizerBuilder.
+
+References:
+  - Nubank nuFormer: sign(2) + amount_bucket(21) + calendar(74) tokenization
+  - TP-BERTa (arXiv:2403.01841): Relative Magnitude Tokenization for numbers
+  - Banking TF (arXiv:2410.08243): date + amount + wording composite tokens
+  - Temporal Tokenization (arXiv:2512.13618): log-based bins for skewed financial data
+"""
+
+import json
+import math
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
+
+class BaseFieldTokenizer:
+    """Base class for all field tokenizers."""
+
+    def __init__(self, prefix: str):
+        self.prefix = prefix
+
+    @property
+    def vocab(self) -> List[str]:
+        """All possible token strings this tokenizer can produce."""
+        raise NotImplementedError
+
+    def __call__(self, value: Any) -> Union[str, List[str]]:
+        """Tokenize a single value. Returns one token string or a list."""
+        raise NotImplementedError
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.vocab)
+
+    def to_dict(self) -> Dict:
+        """Serialize tokenizer state for saving."""
+        return {"type": self.__class__.__name__, "prefix": self.prefix}
+
+    @classmethod
+    def from_dict(cls, d: Dict) -> "BaseFieldTokenizer":
+        """Deserialize tokenizer state."""
+        raise NotImplementedError
+
+
+class SignTokenizer(BaseFieldTokenizer):
+    """Tokenizes the sign of a numerical value.
+    
+    Nubank uses this for credit/debit distinction (2 tokens).
+    Can be generalized to inflow/outflow, buy/sell, etc.
+    
+    Example:
+        >>> tok = SignTokenizer("AMT_SIGN")
+        >>> tok(79.99)   # -> "[AMT_SIGN_POS]"
+        >>> tok(-50.0)   # -> "[AMT_SIGN_NEG]"
+    """
+
+    def __init__(self, prefix: str = "SIGN", pos_label: str = "POS", neg_label: str = "NEG"):
+        super().__init__(prefix)
+        self.pos_label = pos_label
+        self.neg_label = neg_label
+        self._pos_token = f"[{prefix}_{pos_label}]"
+        self._neg_token = f"[{prefix}_{neg_label}]"
+
+    @property
+    def vocab(self) -> List[str]:
+        return [self._pos_token, self._neg_token]
+
+    def __call__(self, value: float) -> str:
+        if value is None or (isinstance(value, float) and math.isnan(value)):
+            return self._pos_token  # default to positive for missing
+        return self._pos_token if value >= 0 else self._neg_token
+
+    def to_dict(self) -> Dict:
+        return {**super().to_dict(), "pos_label": self.pos_label, "neg_label": self.neg_label}
+
+
+class MagnitudeBucketTokenizer(BaseFieldTokenizer):
+    """Quantizes continuous values into bins using quantile-based binning.
+    
+    Follows Nubank's 21-bin quantization and TP-BERTa's Relative Magnitude
+    Tokenization principle. Uses absolute values so sign and magnitude are
+    tokenized independently.
+    
+    Must be fit on training data before use.
+    
+    Example:
+        >>> tok = MagnitudeBucketTokenizer("AMT", n_bins=21)
+        >>> tok.fit(np.array([1.0, 5.0, 10.0, 50.0, 100.0, 500.0]))
+        >>> tok(79.99)   # -> "[AMT_15]" (some bin in the upper range)
+    """
+
+    def __init__(self, prefix: str = "AMT", n_bins: int = 21):
+        super().__init__(prefix)
+        self.n_bins = n_bins
+        self.bin_edges: Optional[np.ndarray] = None
+        self._is_fitted = False
+
+    @property
+    def vocab(self) -> List[str]:
+        return [f"[{self.prefix}_{i:02d}]" for i in range(self.n_bins)]
+
+    def fit(self, values: np.ndarray) -> "MagnitudeBucketTokenizer":
+        """Compute bin edges from training data using quantiles on absolute values."""
+        values = np.asarray(values, dtype=np.float64)
+        # Filter NaN and take absolute values
+        valid = values[~np.isnan(values)]
+        abs_vals = np.abs(valid)
+
+        if len(abs_vals) == 0:
+            raise ValueError("Cannot fit on empty array")
+
+        # Compute quantile edges
+        quantiles = np.linspace(0, 100, self.n_bins + 1)
+        self.bin_edges = np.unique(np.percentile(abs_vals, quantiles))
+
+        # If too few unique edges (degenerate distribution), use linspace
+        if len(self.bin_edges) < 3:
+            self.bin_edges = np.linspace(abs_vals.min(), abs_vals.max(), self.n_bins + 1)
+
+        self._is_fitted = True
+        return self
+
+    def __call__(self, value: float) -> str:
+        if not self._is_fitted:
+            raise RuntimeError(f"MagnitudeBucketTokenizer({self.prefix}) not fitted. Call .fit() first.")
+
+        if value is None or (isinstance(value, float) and math.isnan(value)):
+            return f"[{self.prefix}_00]"  # default to lowest bin for missing
+
+        abs_val = abs(float(value))
+        # searchsorted on interior edges (exclude first and last)
+        bin_idx = int(np.searchsorted(self.bin_edges[1:-1], abs_val))
+        # Clamp to valid range
+        bin_idx = min(bin_idx, self.n_bins - 1)
+        return f"[{self.prefix}_{bin_idx:02d}]"
+
+    def to_dict(self) -> Dict:
+        d = {**super().to_dict(), "n_bins": self.n_bins, "is_fitted": self._is_fitted}
+        if self._is_fitted:
+            d["bin_edges"] = self.bin_edges.tolist()
+        return d
+
+    @classmethod
+    def from_dict(cls, d: Dict) -> "MagnitudeBucketTokenizer":
+        tok = cls(prefix=d["prefix"], n_bins=d["n_bins"])
+        if d.get("is_fitted") and "bin_edges" in d:
+            tok.bin_edges = np.array(d["bin_edges"])
+            tok._is_fitted = True
+        return tok
+
+
+class DiscreteNumericalTokenizer(BaseFieldTokenizer):
+    """Tokenizes small discrete numerical values (quantities, counts).
+    
+    Maps integers 0..max_value to individual tokens, with an overflow token
+    for values exceeding max_value.
+    
+    Example:
+        >>> tok = DiscreteNumericalTokenizer("QTY", max_value=10)
+        >>> tok(3)    # -> "[QTY_03]"
+        >>> tok(15)   # -> "[QTY_OVER]"
+    """
+
+    def __init__(self, prefix: str = "QTY", max_value: int = 10):
+        super().__init__(prefix)
+        self.max_value = max_value
+        self._overflow_token = f"[{prefix}_OVER]"
+
+    @property
+    def vocab(self) -> List[str]:
+        tokens = [f"[{self.prefix}_{i:02d}]" for i in range(self.max_value + 1)]
+        tokens.append(self._overflow_token)
+        return tokens
+
+    def __call__(self, value: Any) -> str:
+        if value is None:
+            return f"[{self.prefix}_00]"
+        v = int(value)
+        if v < 0:
+            v = 0
+        if v > self.max_value:
+            return self._overflow_token
+        return f"[{self.prefix}_{v:02d}]"
+
+    def to_dict(self) -> Dict:
+        return {**super().to_dict(), "max_value": self.max_value}
+
+
+class CalendarTokenizer(BaseFieldTokenizer):
+    """Decomposes timestamps into calendar component tokens.
+    
+    Follows Nubank's approach: month(12) + dow(7) + dom(31) + hour(24) = 74 tokens.
+    Accepts datetime objects or ISO format strings.
+    
+    Example:
+        >>> tok = CalendarTokenizer("TS", fields=["month", "dow", "dom", "hour"])
+        >>> tok(datetime(2025, 3, 15, 14, 30))
+        ['[TS_MON_03]', '[TS_DOW_5]', '[TS_DOM_15]', '[TS_HOUR_14]']
+    """
+
+    # Maps field name -> (token format, extraction function, count)
+    FIELD_REGISTRY = {
+        "month":      (lambda p, i: f"[{p}_MON_{i+1:02d}]",  lambda dt: dt.month - 1,  12),
+        "dow":        (lambda p, i: f"[{p}_DOW_{i}]",         lambda dt: dt.weekday(),   7),
+        "dom":        (lambda p, i: f"[{p}_DOM_{i+1:02d}]",   lambda dt: dt.day - 1,     31),
+        "hour":       (lambda p, i: f"[{p}_HOUR_{i:02d}]",    lambda dt: dt.hour,        24),
+        "quarter":    (lambda p, i: f"[{p}_Q{i+1}]",          lambda dt: (dt.month-1)//3, 4),
+        "minute_bin": (lambda p, i: f"[{p}_MINBIN_{i}]",      lambda dt: dt.minute // 15, 4),
+    }
+
+    def __init__(self, prefix: str = "TS", fields: Optional[List[str]] = None):
+        super().__init__(prefix)
+        self.fields = fields or ["month", "dow", "dom", "hour"]
+        # Validate
+        for f in self.fields:
+            if f not in self.FIELD_REGISTRY:
+                raise ValueError(f"Unknown calendar field: '{f}'. Available: {list(self.FIELD_REGISTRY.keys())}")
+
+    @property
+    def vocab(self) -> List[str]:
+        tokens = []
+        for field_name in self.fields:
+            fmt_fn, _, count = self.FIELD_REGISTRY[field_name]
+            tokens.extend(fmt_fn(self.prefix, i) for i in range(count))
+        return tokens
+
+    def _parse_datetime(self, value: Any) -> datetime:
+        if isinstance(value, datetime):
+            return value
+        if isinstance(value, str):
+            # Try common formats
+            for fmt in ("%Y-%m-%dT%H:%M:%S", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
+                try:
+                    return datetime.strptime(value, fmt)
+                except ValueError:
+                    continue
+            raise ValueError(f"Cannot parse datetime string: {value}")
+        raise TypeError(f"Expected datetime or str, got {type(value)}")
+
+    def __call__(self, value: Any) -> List[str]:
+        dt = self._parse_datetime(value)
+        tokens = []
+        for field_name in self.fields:
+            fmt_fn, extract_fn, count = self.FIELD_REGISTRY[field_name]
+            idx = extract_fn(dt)
+            idx = max(0, min(idx, count - 1))  # clamp
+            tokens.append(fmt_fn(self.prefix, idx))
+        return tokens
+
+    def to_dict(self) -> Dict:
+        return {**super().to_dict(), "fields": self.fields}
+
+
+class CategoricalTokenizer(BaseFieldTokenizer):
+    """Maps categorical string values to fixed vocabulary tokens.
+    
+    Unknown values map to an [PREFIX_UNK] token.
+    
+    Example:
+        >>> tok = CategoricalTokenizer("EVT", ["view", "purchase", "return"])
+        >>> tok("purchase")    # -> "[EVT_001]"
+        >>> tok("unknown")     # -> "[EVT_UNK]"
+    """
+
+    def __init__(self, prefix: str, categories: List[str]):
+        super().__init__(prefix)
+        self.categories = list(categories)
+        self._token_map = {cat: f"[{prefix}_{i:03d}]" for i, cat in enumerate(categories)}
+        self._unk_token = f"[{prefix}_UNK]"
+        # Also build reverse map for decoding
+        self._reverse_map = {v: k for k, v in self._token_map.items()}
+        self._reverse_map[self._unk_token] = "<unknown>"
+
+    @property
+    def vocab(self) -> List[str]:
+        return list(self._token_map.values()) + [self._unk_token]
+
+    def __call__(self, value: Any) -> str:
+        if value is None:
+            return self._unk_token
+        return self._token_map.get(str(value), self._unk_token)
+
+    def decode_token(self, token: str) -> str:
+        """Map a token string back to its category value."""
+        return self._reverse_map.get(token, "<unknown>")
+
+    def to_dict(self) -> Dict:
+        return {**super().to_dict(), "categories": self.categories}
+
+    @classmethod
+    def from_dict(cls, d: Dict) -> "CategoricalTokenizer":
+        return cls(prefix=d["prefix"], categories=d["categories"])
+
+
+# =============================================================================
+# Factory function to create field tokenizer from FieldSpec
+# =============================================================================
+
+def create_field_tokenizer(spec) -> BaseFieldTokenizer:
+    """Create the appropriate field tokenizer from a FieldSpec.
+    
+    Args:
+        spec: A FieldSpec instance from schema.py
+        
+    Returns:
+        An initialized BaseFieldTokenizer subclass
+    """
+    from ..schema import FieldType  # avoid circular import
+    
+    if spec.field_type == FieldType.SIGN:
+        return SignTokenizer(prefix=spec.prefix)
+    
+    elif spec.field_type == FieldType.NUMERICAL_CONTINUOUS:
+        return MagnitudeBucketTokenizer(prefix=spec.prefix, n_bins=spec.n_bins)
+    
+    elif spec.field_type == FieldType.NUMERICAL_DISCRETE:
+        return DiscreteNumericalTokenizer(prefix=spec.prefix, max_value=spec.max_value)
+    
+    elif spec.field_type == FieldType.CATEGORICAL_FIXED:
+        return CategoricalTokenizer(prefix=spec.prefix, categories=spec.categories)
+    
+    elif spec.field_type == FieldType.TEMPORAL:
+        return CalendarTokenizer(prefix=spec.prefix, fields=spec.calendar_fields)
+    
+    elif spec.field_type == FieldType.TEXT:
+        return None  # Text is handled by the BPE tokenizer in DomainTokenizerBuilder
+    
+    else:
+        raise ValueError(f"Unknown field type: {spec.field_type}")