morphbpe-pl 0.3.0__py3-none-any.whl

morphbpe/tokenizer.py

@@ -0,0 +1,614 @@
+"""MorphBPE tokenizer: BPE with morphological boundary penalties."""
+
+import collections
+import json
+import re
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Set, Tuple, Union
+
+import numpy as np
+from tqdm import tqdm
+
+from morphbpe.morfeusz_backend import MorfeuszBackend
+
+# ── Special tokens ──────────────────────────────────────────────────
+
+SPECIAL_TOKENS = {
+    "<pad>": 0,
+    "<unk>": 1,
+    "<cls>": 2,
+    "<sep>": 3,
+    "<mask>": 4,
+}
+
+# ── Byte-level tokens ──────────────────────────────────────────────
+# 256 byte tokens for fallback encoding of any UTF-8 input.
+# IDs 5..260 are reserved for bytes. Character vocab starts at 261.
+
+def _make_byte_tokens() -> Tuple[Dict[str, int], Dict[int, str]]:
+    """Create byte-level token mappings. Returns (token2id, id2token)."""
+    t2i: Dict[str, int] = {}
+    i2t: Dict[int, str] = {}
+    for i in range(256):
+        # Use printable representation: raw byte for printable ASCII,
+        # hex escape for others (similar to GPT-2's byte encoder)
+        if 33 <= i <= 126 and i != 96:  # printable except backtick
+            ch = chr(i)
+        else:
+            ch = f"<0x{i:02X}>"
+        t2i[ch] = i + 5
+        i2t[i + 5] = ch
+    return t2i, i2t
+
+BYTE_TOKENS, BYTE_TOKENS_INV = _make_byte_tokens()
+
+# Character that maps to each byte value for UTF-8 fallback
+_BYTE_CHARS: Dict[int, str] = {}
+for _i in range(256):
+    if 33 <= _i <= 126 and _i != 96:
+        _BYTE_CHARS[_i] = chr(_i)
+    else:
+        _BYTE_CHARS[_i] = f"<0x{_i:02X}>"
+
+# ID where regular character vocab starts (after special + byte tokens)
+VOCAB_OFFSET = 261
+
+# ── Pre-tokenization regex ──────────────────────────────────────────
+# Splits punctuation, contractions, numbers, and whitespace.
+# Similar to GPT-2's tokenizer but adapted for Polish.
+
+_PRE_TOKENIZE_RE = re.compile(
+    r"""
+    (?:[^\W\d_]+)          |  # letters (including Polish diacritics)
+    (?:\d[\d\.]*\d)        |  # numbers (including decimals)
+    (?:\d+)                |  # single digits
+    (?:[^\w\s])            |  # punctuation (single char)
+    (?:\s+)                   # whitespace
+    """,
+    re.UNICODE | re.VERBOSE,
+)
+
+
+def pre_tokenize(text: str) -> List[str]:
+    """Split text into pre-tokens (words, punctuation, numbers)."""
+    return _PRE_TOKENIZE_RE.findall(text)
+
+
+# ── Trie for O(m) encoding ──────────────────────────────────────────
+
+
+class _TrieNode:
+    """Trie node for fast token lookup."""
+
+    __slots__ = ("children", "token_id")
+
+    def __init__(self):
+        self.children: Dict[str, "_TrieNode"] = {}
+        self.token_id: Optional[int] = None
+
+
+class _Trie:
+    """Trie for longest-match-first tokenization."""
+
+    def __init__(self):
+        self.root = _TrieNode()
+
+    def insert(self, token: str, token_id: int) -> None:
+        node = self.root
+        for ch in token:
+            if ch not in node.children:
+                node.children[ch] = _TrieNode()
+            node = node.children[ch]
+        node.token_id = token_id
+
+    def tokenize(self, word: str, unk_id: int) -> List[Tuple[str, int]]:
+        """Greedy longest-match tokenization with byte-level fallback.
+
+        Returns (token_str, id) pairs. For unknown characters, encodes
+        as UTF-8 bytes using the byte token table.
+        """
+        result: List[Tuple[str, int]] = []
+        i = 0
+        n = len(word)
+        while i < n:
+            node = self.root
+            best_end = -1
+            best_id = -1
+            j = i
+            while j < n and word[j] in node.children:
+                node = node.children[word[j]]
+                j += 1
+                if node.token_id is not None:
+                    best_end = j
+                    best_id = node.token_id
+            if best_end == -1:
+                # No match — encode character as UTF-8 bytes
+                ch = word[i]
+                ch_bytes = ch.encode("utf-8")
+                for b in ch_bytes:
+                    byte_str = _BYTE_CHARS[b]
+                    byte_id = BYTE_TOKENS[byte_str]
+                    result.append((byte_str, byte_id))
+                i += 1
+            else:
+                result.append((word[i:best_end], best_id))
+                i = best_end
+        return result
+
+
+# ── MorphBPE Tokenizer ──────────────────────────────────────────────
+
+
+class MorphBPETokenizer:
+    """Morphologically-aware BPE tokenizer.
+
+    Parameters
+    ----------
+    vocab_size : int
+        Target vocabulary size (default 1000).
+    alpha : float
+        Penalty strength for boundary-crossing merges.
+        0.0 = standard BPE, 1.0 = full block, >1.0 = over-penalization.
+    constraint : str
+        Which boundaries to enforce: "suffix", "prefix", or "both".
+    use_morfeusz : bool
+        Use Morfeusz2 if available; otherwise heuristic fallback.
+    """
+
+    def __init__(
+        self,
+        vocab_size: int = 1000,
+        alpha: float = 1.0,
+        constraint: str = "suffix",
+        use_morfeusz: bool = True,
+    ):
+        if constraint not in ("suffix", "prefix", "both"):
+            raise ValueError(f"constraint must be suffix/prefix/both, got {constraint}")
+
+        self.vocab_size = vocab_size
+        self.alpha = alpha
+        self.constraint = constraint
+
+        self._backend = MorfeuszBackend(use_morfeusz=use_morfeusz)
+        self.token2id: Dict[str, int] = {}
+        self.id2token: Dict[int, str] = {}
+        self.merges: List[Tuple[str, str]] = []
+
+        # Special tokens are always present
+        self._special_tokens = dict(SPECIAL_TOKENS)
+        self._inv_special_tokens: Dict[int, str] = {v: k for k, v in SPECIAL_TOKENS.items()}
+        self.pad_token_id = 0
+        self.unk_token_id = 1
+        self.cls_token_id = 2
+        self.sep_token_id = 3
+        self.mask_token_id = 4
+
+        # Trie for fast encoding (built after training/loading)
+        self._trie: Optional[_Trie] = None
+        self._encode_cache: Dict[str, List[str]] = {}
+
+    # ── Properties ──────────────────────────────────────────────────
+
+    @property
+    def special_tokens(self) -> Dict[str, int]:
+        return dict(self._special_tokens)
+
+    @property
+    def all_special_tokens(self) -> List[str]:
+        return list(self._special_tokens.keys())
+
+    def __len__(self) -> int:
+        return len(self.token2id)
+
+    def __contains__(self, token: str) -> bool:
+        return token in self.token2id
+
+    # ── Trie management ─────────────────────────────────────────────
+
+    def _build_trie(self) -> None:
+        """Build trie from current vocab for fast encoding."""
+        self._trie = _Trie()
+        for token, idx in self.token2id.items():
+            self._trie.insert(token, idx)
+        self._encode_cache.clear()
+
+    # ── Training ────────────────────────────────────────────────────
+
+    def train(self, texts: List[str], verbose: bool = True) -> float:
+        """Train tokenizer on a list of texts. Returns training time in seconds."""
+        start = time.time()
+
+        word_freq: Dict[str, int] = collections.defaultdict(int)
+        boundaries: Dict[str, Tuple[Set[int], Set[int]]] = {}
+        it = tqdm(texts, desc="Preprocessing") if verbose else texts
+        for text in it:
+            # Use pre-tokenizer instead of naive split
+            tokens = pre_tokenize(text)
+            for tok in tokens:
+                if tok.isspace():
+                    continue  # skip whitespace tokens
+                word_freq[" ".join(list(tok))] += 1
+                if tok not in boundaries:
+                    boundaries[tok] = self._backend.get_boundaries(tok)
+
+        # Initial character vocab (after special tokens)
+        chars: Set[str] = set()
+        for w in word_freq:
+            chars.update(w.split())
+
+        # Initialize vocab with special tokens + byte tokens
+        self.token2id = dict(SPECIAL_TOKENS)
+        self.token2id.update(BYTE_TOKENS)
+        self.id2token = {v: k for k, v in self.token2id.items()}
+
+        # Character vocab starts after special + byte tokens
+        next_id = VOCAB_OFFSET
+
+        for c in sorted(chars):
+            if c not in self.token2id:
+                self.token2id[c] = next_id
+                self.id2token[next_id] = c
+                next_id += 1
+
+        n_merges = self.vocab_size - len(self.token2id)
+        it = tqdm(range(n_merges), desc="Merging") if verbose else range(n_merges)
+        for _ in it:
+            pairs = self._pair_scores(word_freq, boundaries)
+            if not pairs:
+                break
+            best = max(pairs, key=pairs.get)
+            self.merges.append(best)
+            word_freq = self._apply_merge(word_freq, best)
+            token = "".join(best)
+            if token not in self.token2id:
+                self.token2id[token] = next_id
+                self.id2token[next_id] = token
+                next_id += 1
+
+        # Build trie for fast encoding
+        self._build_trie()
+
+        elapsed = time.time() - start
+        if verbose:
+            print(
+                f"Trained: vocab={len(self.token2id)}, "
+                f"merges={len(self.merges)}, time={elapsed:.1f}s"
+            )
+        return elapsed
+
+    def train_from_file(self, path: str, verbose: bool = True) -> float:
+        """Train from a plain-text file (one document per line)."""
+        texts = Path(path).read_text(encoding="utf-8").splitlines()
+        texts = [t for t in texts if t.strip()]
+        return self.train(texts, verbose=verbose)
+
+    # ── Encode / Decode ─────────────────────────────────────────────
+
+    def encode_word(self, word: str) -> List[str]:
+        """Tokenize a single word into subword strings."""
+        if not self._trie:
+            self._build_trie()
+
+        if word in self._encode_cache:
+            return self._encode_cache[word]
+
+        # Use trie for fast encoding
+        pairs = self._trie.tokenize(word, self.unk_token_id)
+        result = [tok_str for tok_str, _ in pairs]
+        self._encode_cache[word] = result
+        return result
+
+    def encode(
+        self,
+        text: str,
+        add_special_tokens: bool = False,
+    ) -> List[int]:
+        """Encode text to a list of token IDs.
+
+        Parameters
+        ----------
+        text : str
+            Input text.
+        add_special_tokens : bool
+            If True, prepend [CLS] and append [SEP].
+
+        Returns
+        -------
+        List[int]
+            Token IDs.
+        """
+        if not self._trie:
+            self._build_trie()
+
+        ids: List[int] = []
+        if add_special_tokens:
+            ids.append(self.cls_token_id)
+
+        # Pre-tokenize to handle punctuation
+        tokens = pre_tokenize(text)
+        for tok in tokens:
+            if tok.isspace():
+                continue  # skip whitespace
+            for subtok in self.encode_word(tok):
+                ids.append(self.token2id.get(subtok, self.unk_token_id))
+
+        if add_special_tokens:
+            ids.append(self.sep_token_id)
+
+        return ids
+
+    def decode(
+        self,
+        ids: List[int],
+        skip_special_tokens: bool = True,
+        spaces_between_words: bool = True,
+    ) -> str:
+        """Decode token IDs back to text.
+
+        Parameters
+        ----------
+        ids : List[int]
+            Token IDs to decode.
+        skip_special_tokens : bool
+            If True, skip special tokens in output.
+        spaces_between_words : bool
+            If True, insert spaces between word-initial tokens.
+
+        Returns
+        -------
+        str
+            Decoded text.
+        """
+        special_ids = set(SPECIAL_TOKENS.values())
+        parts: List[str] = []
+        for i in ids:
+            if skip_special_tokens and i in special_ids:
+                continue
+            token = self.id2token.get(i, "<unk>")
+            if token == "<unk>":
+                continue
+            parts.append(token)
+
+        if not spaces_between_words:
+            return "".join(parts)
+
+        # Heuristic: insert space before tokens that look like word starts
+        # (lowercase letter or Polish uppercase at word boundary)
+        result = parts[0] if parts else ""
+        for i in range(1, len(parts)):
+            prev = parts[i - 1]
+            curr = parts[i]
+            # Space if current token starts a new word:
+            # - previous token ended with punctuation
+            # - current token is a standalone character that's likely a word start
+            if (
+                prev[-1] in ".,;:!?)]}\u201d\u2019"
+                or curr[0] in "({[\u201c\u2018"
+                or (len(prev) == 1 and prev in "iyIY")  # Polish conjunctions
+            ):
+                result += " " + curr
+            else:
+                result += curr
+        return result
+
+    def batch_encode(
+        self,
+        texts: List[str],
+        add_special_tokens: bool = False,
+        padding: bool = True,
+        max_length: Optional[int] = None,
+        return_attention_mask: bool = True,
+    ) -> Dict[str, List[List[int]]]:
+        """Encode a batch of texts.
+
+        Parameters
+        ----------
+        texts : List[str]
+            Input texts.
+        add_special_tokens : bool
+            If True, prepend [CLS] and append [SEP].
+        padding : bool
+            If True, pad to max length in batch.
+        max_length : int, optional
+            Maximum sequence length (truncates if longer).
+        return_attention_mask : bool
+            If True, return attention masks.
+
+        Returns
+        -------
+        dict
+            {"input_ids": List[List[int]], "attention_mask": List[List[int]]}
+        """
+        batch_ids = [self.encode(t, add_special_tokens=add_special_tokens) for t in texts]
+
+        if max_length:
+            batch_ids = [ids[:max_length] for ids in batch_ids]
+
+        result = {"input_ids": batch_ids}
+
+        if padding:
+            max_len = max(len(ids) for ids in batch_ids) if batch_ids else 0
+            padded = []
+            masks = []
+            for ids in batch_ids:
+                mask = [1] * len(ids) + [0] * (max_len - len(ids))
+                padded.append(ids + [self.pad_token_id] * (max_len - len(ids)))
+                masks.append(mask)
+            result["input_ids"] = padded
+            if return_attention_mask:
+                result["attention_mask"] = masks
+        elif return_attention_mask:
+            result["attention_mask"] = [[1] * len(ids) for ids in batch_ids]
+
+        return result
+
+    @property
+    def vocab(self) -> Dict[str, int]:
+        """HuggingFace-compatible vocab property."""
+        return dict(self.token2id)
+
+    def get_vocab_size(self) -> int:
+        return len(self.token2id)
+
+    # ── Save / Load ─────────────────────────────────────────────────
+
+    def save(self, directory: str) -> None:
+        """Save tokenizer to a directory."""
+        path = Path(directory)
+        path.mkdir(parents=True, exist_ok=True)
+
+        config = {
+            "vocab_size": self.vocab_size,
+            "alpha": self.alpha,
+            "constraint": self.constraint,
+            "has_morfeusz": self._backend.has_morfeusz,
+            "special_tokens": self._special_tokens,
+            "model_type": "morphbpe",
+        }
+        (path / "config.json").write_text(
+            json.dumps(config, indent=2, ensure_ascii=False), encoding="utf-8"
+        )
+        (path / "vocab.json").write_text(
+            json.dumps(self.token2id, indent=2, ensure_ascii=False), encoding="utf-8"
+        )
+        merges_lines = [f"{a} {b}" for a, b in self.merges]
+        (path / "merges.txt").write_text(
+            "\n".join(merges_lines), encoding="utf-8"
+        )
+
+    @classmethod
+    def load(cls, directory: str) -> "MorphBPETokenizer":
+        """Load tokenizer from a directory."""
+        path = Path(directory)
+        config = json.loads((path / "config.json").read_text(encoding="utf-8"))
+        tok = cls(
+            vocab_size=config["vocab_size"],
+            alpha=config["alpha"],
+            constraint=config["constraint"],
+            use_morfeusz=config.get("has_morfeusz", True),
+        )
+        tok.token2id = json.loads((path / "vocab.json").read_text(encoding="utf-8"))
+        tok.id2token = {}
+        for k, v in tok.token2id.items():
+            tok.id2token[int(v)] = k
+        raw = (path / "merges.txt").read_text(encoding="utf-8").strip()
+        if raw:
+            tok.merges = [tuple(line.split(" ", 1)) for line in raw.splitlines()]
+
+        # Load special tokens if present
+        if "special_tokens" in config:
+            tok._special_tokens = config["special_tokens"]
+            tok._inv_special_tokens = {v: k for k, v in tok._special_tokens.items()}
+
+        # Build trie for fast encoding
+        tok._build_trie()
+
+        return tok
+
+    # ── Metrics ─────────────────────────────────────────────────────
+
+    def fertility(self, texts: List[str]) -> float:
+        """Average tokens per word."""
+        n_tok = n_word = 0
+        for text in texts:
+            for word in text.split():
+                n_tok += len(self.encode_word(word))
+                n_word += 1
+        return n_tok / n_word if n_word else 0.0
+
+    def morph_edit_distance(self, texts: List[str]) -> float:
+        """Mean MorphED across words (≥3 chars)."""
+        vals = self._per_word_morphed(texts)
+        return float(np.mean(vals)) if vals else 0.0
+
+    def morph_edit_distance_ci(
+        self, texts: List[str], n_resamples: int = 1000, seed: int = 42
+    ) -> Tuple[float, float, float]:
+        """MorphED with bootstrap 95% CI → (mean, lower, upper)."""
+        vals = np.array(self._per_word_morphed(texts))
+        if len(vals) == 0:
+            return 0.0, 0.0, 0.0
+        rng = np.random.RandomState(seed)
+        means = [
+            np.mean(rng.choice(vals, size=len(vals), replace=True))
+            for _ in range(n_resamples)
+        ]
+        return float(np.mean(vals)), float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5))
+
+    def _per_word_morphed(self, texts: List[str]) -> List[int]:
+        result: List[int] = []
+        for text in texts:
+            for word in text.split():
+                if len(word) < 3:
+                    continue
+                prefix_b, suffix_b = self._backend.get_boundaries(word)
+                all_b = sorted(prefix_b | suffix_b)
+                morphemes: List[str] = []
+                prev = 0
+                for b in all_b:
+                    morphemes.append(word[prev:b])
+                    prev = b
+                morphemes.append(word[prev:])
+                tokens = self.encode_word(word)
+                result.append(_list_edit_distance(morphemes, tokens))
+        return result
+
+    # ── Internal BPE logic ──────────────────────────────────────────
+
+    def _active_boundaries(
+        self, prefix_b: Set[int], suffix_b: Set[int]
+    ) -> Set[int]:
+        if self.constraint == "prefix":
+            return prefix_b
+        if self.constraint == "suffix":
+            return suffix_b
+        return prefix_b | suffix_b
+
+    def _pair_scores(
+        self,
+        word_freq: Dict[str, int],
+        boundaries: Dict[str, Tuple[Set[int], Set[int]]],
+    ) -> Dict[Tuple[str, str], float]:
+        pairs: Dict[Tuple[str, str], float] = collections.defaultdict(float)
+        for word, freq in word_freq.items():
+            tokens = word.split()
+            orig = word.replace(" ", "")
+            pfx_b, sfx_b = boundaries.get(orig, (set(), set()))
+            active = self._active_boundaries(pfx_b, sfx_b)
+
+            pos = 0
+            for i in range(len(tokens) - 1):
+                tlen = len(tokens[i])
+                crosses = (pos + tlen) in active
+                if crosses:
+                    pairs[(tokens[i], tokens[i + 1])] += freq * (1.0 - self.alpha)
+                else:
+                    pairs[(tokens[i], tokens[i + 1])] += freq
+                pos += tlen
+        return pairs
+
+    @staticmethod
+    def _apply_merge(
+        word_freq: Dict[str, int], pair: Tuple[str, str]
+    ) -> Dict[str, int]:
+        bigram = " ".join(pair)
+        repl = "".join(pair)
+        return {w.replace(bigram, repl): f for w, f in word_freq.items()}
+
+
+def _list_edit_distance(a: List[str], b: List[str]) -> int:
+    """Edit distance between two lists of strings."""
+    m, n = len(a), len(b)
+    dp = [[0] * (n + 1) for _ in range(m + 1)]
+    for i in range(m + 1):
+        dp[i][0] = i
+    for j in range(n + 1):
+        dp[0][j] = j
+    for i in range(1, m + 1):
+        for j in range(1, n + 1):
+            if a[i - 1] == b[j - 1]:
+                dp[i][j] = dp[i - 1][j - 1]
+            else:
+                dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])
+    return dp[m][n]