""" Domain Tokenizer Builder — assembles per-field tokenizers into an HF-compatible PreTrainedTokenizerFast. This is the core of domainTokenizer: it takes a DomainSchema, builds per-field tokenizers, fits data-dependent ones, and produces a single HuggingFace tokenizer that can encode domain events as token ID sequences. The output tokenizer is fully compatible with HF Trainer, push_to_hub, from_pretrained, etc. IMPORTANT: Uses Whitespace pre-tokenizer (not ByteLevel) because our sequences are space-separated special tokens like '[EVT_000] [PRICE_16] electronics.smartphone'. ByteLevel would split these into byte-fragments causing massive UNK rates. """ import json import os from datetime import datetime from pathlib import Path from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union import numpy as np from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers from transformers import PreTrainedTokenizerFast from ..schema import DomainSchema, FieldSpec, FieldType from .field_tokenizers import ( BaseFieldTokenizer, CalendarTokenizer, CategoricalTokenizer, DiscreteNumericalTokenizer, MagnitudeBucketTokenizer, SignTokenizer, create_field_tokenizer, ) # Control tokens -- always present CONTROL_TOKENS = ["[PAD]", "[UNK]", "[BOS]", "[EOS]", "[MASK]", "[CLS]", "[SEP]"] class DomainTokenizerBuilder: """Builds an HF-compatible tokenizer from a DomainSchema. Workflow: 1. builder = DomainTokenizerBuilder(schema) 2. builder.fit(events) # fit magnitude bins etc. 3. hf_tok = builder.build(text_corpus) # build HF tokenizer 4. tokens = builder.tokenize_event(event) # tokenize a single event 5. ids = hf_tok(tokens_str) # convert to IDs Or use the convenience method: 6. ids = builder.encode_event(event, hf_tok) # event -> IDs in one call 7. ids = builder.encode_sequence(events, hf_tok) # full sequence -> IDs """ def __init__(self, schema: DomainSchema): self.schema = schema self.field_tokenizers: Dict[str, Optional[BaseFieldTokenizer]] = {} self._is_fitted = False self._build_field_tokenizers() def _build_field_tokenizers(self): for spec in self.schema.fields: self.field_tokenizers[spec.name] = create_field_tokenizer(spec) def fit(self, events: Sequence[Dict[str, Any]]) -> "DomainTokenizerBuilder": """Fit data-dependent tokenizers on training events.""" for spec in self.schema.fields: if spec.field_type == FieldType.NUMERICAL_CONTINUOUS: tok = self.field_tokenizers[spec.name] values = [] for event in events: v = event.get(spec.name) if isinstance(event, dict) else getattr(event, spec.name, None) if v is not None: values.append(float(v)) if values: tok.fit(np.array(values)) else: raise ValueError(f"No values found for field '{spec.name}' during fitting") self._is_fitted = True return self @property def is_fitted(self) -> bool: if not self.schema.fittable_field_names: return True return self._is_fitted def _collect_special_tokens(self) -> List[str]: """Collect all special tokens: control + event separator + per-field domain tokens.""" tokens = list(CONTROL_TOKENS) tokens.append(self.schema.event_separator) for spec in self.schema.fields: tok = self.field_tokenizers.get(spec.name) if tok is not None and hasattr(tok, "vocab"): tokens.extend(tok.vocab) seen = set() unique = [] for t in tokens: if t not in seen: seen.add(t) unique.append(t) return unique def build( self, text_corpus: Optional[Iterator[str]] = None, bpe_vocab_size: int = 8000, min_frequency: int = 2, ) -> PreTrainedTokenizerFast: """Build a complete HuggingFace-compatible tokenizer. Uses Whitespace pre-tokenizer (not ByteLevel) because domain token sequences are space-separated: '[EVT_000] [PRICE_16] electronics.smartphone'. Whitespace splits on spaces, keeping special tokens intact. BPE then handles subword splitting within text tokens (e.g., 'electronics.smartphone' -> 'electronics', '.', 'smartphone'). """ for name in self.schema.fittable_field_names: tok = self.field_tokenizers[name] if isinstance(tok, MagnitudeBucketTokenizer) and not tok._is_fitted: raise RuntimeError( f"Field '{name}' requires fitting. Call builder.fit(events) first." ) all_special_tokens = self._collect_special_tokens() # Build BPE tokenizer with Whitespace pre-tokenizer # Whitespace splits on spaces → each special token like [EVT_000] stays intact # BPE merges then handle subword splitting within text fields base_tokenizer = Tokenizer(models.BPE(unk_token="[UNK]")) base_tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() base_tokenizer.decoder = decoders.BPEDecoder() if self.schema.has_text_fields and text_corpus is not None: trainer_obj = trainers.BpeTrainer( vocab_size=bpe_vocab_size, special_tokens=all_special_tokens, min_frequency=min_frequency, show_progress=True, ) if isinstance(text_corpus, (list, tuple)): base_tokenizer.train_from_iterator(iter(text_corpus), trainer=trainer_obj) else: base_tokenizer.train_from_iterator(text_corpus, trainer=trainer_obj) else: # No text fields — vocabulary-only tokenizer vocab = {token: i for i, token in enumerate(all_special_tokens)} merges = [] base_tokenizer = Tokenizer(models.BPE(vocab=vocab, merges=merges, unk_token="[UNK]")) base_tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() base_tokenizer.decoder = decoders.BPEDecoder() hf_tokenizer = PreTrainedTokenizerFast( tokenizer_object=base_tokenizer, bos_token="[BOS]", eos_token="[EOS]", pad_token="[PAD]", unk_token="[UNK]", mask_token="[MASK]", cls_token="[CLS]", sep_token="[SEP]", ) return hf_tokenizer def tokenize_event(self, event: Union[Dict[str, Any], Any]) -> List[str]: """Convert a single domain event into a list of token strings.""" tokens = [] for spec in self.schema.fields: if isinstance(event, dict): value = event.get(spec.name) else: value = getattr(event, spec.name, None) if spec.field_type == FieldType.TEXT: if value is not None: tokens.append(str(value)) continue tok = self.field_tokenizers.get(spec.name) if tok is None: continue if value is None: tokens.append("[UNK]") continue result = tok(value) if isinstance(result, list): tokens.extend(result) else: tokens.append(result) return tokens def tokenize_sequence( self, events: Sequence[Union[Dict[str, Any], Any]], add_bos: bool = True, add_eos: bool = True, ) -> List[str]: """Tokenize a full sequence of events into token strings.""" all_tokens = [] if add_bos: all_tokens.append("[BOS]") for i, event in enumerate(events): if i > 0: all_tokens.append(self.schema.event_separator) event_tokens = self.tokenize_event(event) all_tokens.extend(event_tokens) if add_eos: all_tokens.append("[EOS]") return all_tokens def encode_sequence( self, events: Sequence[Union[Dict[str, Any], Any]], hf_tokenizer: PreTrainedTokenizerFast, max_length: int = 2048, add_bos: bool = True, add_eos: bool = True, return_tensors: Optional[str] = None, ) -> Dict[str, Any]: """Full pipeline: events -> token strings -> token IDs.""" token_strings = self.tokenize_sequence(events, add_bos=add_bos, add_eos=add_eos) token_text = " ".join(token_strings) encoding = hf_tokenizer( token_text, max_length=max_length, truncation=True, padding="max_length", return_tensors=return_tensors, ) return encoding def save(self, directory: str): """Save the builder state (fitted bins, schema, etc.) to a directory.""" os.makedirs(directory, exist_ok=True) state = { "schema_name": self.schema.name, "is_fitted": self._is_fitted, "field_tokenizers": {}, } for name, tok in self.field_tokenizers.items(): if tok is not None: state["field_tokenizers"][name] = tok.to_dict() with open(os.path.join(directory, "domain_tokenizer_builder.json"), "w") as f: json.dump(state, f, indent=2) def get_stats(self) -> Dict[str, Any]: """Return statistics about the tokenizer configuration.""" return { "schema_name": self.schema.name, "total_fields": len(self.schema.fields), "special_token_count": self.schema.special_token_count, "fixed_tokens_per_event": self.schema.fixed_tokens_per_event, "has_text_fields": self.schema.has_text_fields, "is_fitted": self.is_fitted, "field_details": { spec.name: { "type": spec.field_type.value, "vocab_tokens": spec.token_count, "tokens_per_event": spec.tokens_per_event, } for spec in self.schema.fields }, }