File size: 3,044 Bytes
a8e9558
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
"""
BPETokenizer β€” HuggingFace tokenizers wrapper
──────────────────────────────────────────────
Wraps a pre-trained BPE tokenizer loaded from a JSON file.

SPECIAL TOKENS:
  <pad>  β€” Padding token for batching sequences of different lengths.
            Not used during standard next-token-prediction training
            (all windows are the same length), but needed for inference batching.
  <unk>  β€” Unknown token. A well-trained BPE tokenizer should never emit this
            because byte-level BPE can represent any Unicode character.
  <bos>  β€” Beginning-of-Sequence. Prepended to each document during training
            so the model learns to associate it with "start of generation".
  <eos>  β€” End-of-Sequence. Appended to each document. The model is trained
            to predict this token when a natural endpoint is reached. During
            inference, generation stops when EOS is produced.

BYTE-LEVEL BPE:
  The tokenizer uses ByteLevel pre-tokenization (from train_tokenizer.py).
  This means every byte of the input is representable as a single token β€”
  no character is ever truly "unknown". The Δ  prefix (used in GPT-2 / RoBERTa)
  marks word boundaries: a token starting with Δ  follows a space. The clean()
  method converts these back to readable text.
"""

from tokenizers import Tokenizer, decoders


class BPETokenizer:
    """
    Wraps a HuggingFace `Tokenizer` JSON file for easy encode/decode.

    Args:
        tokenizer_path: Path to the saved tokenizer JSON
                        (produced by tokenizer/train_tokenizer.py).
    """

    def __init__(self, tokenizer_path: str = "dataProcessed/tokenizer.json"):
        self.tokenizer = Tokenizer.from_file(tokenizer_path)

        # Ensure ByteLevel decoder is attached so decode() properly maps
        # byte-level tokens (like Ġ for space and Ċ for newline) back to text.
        if self.tokenizer.decoder is None:
            self.tokenizer.decoder = decoders.ByteLevel()

        # Resolve special token IDs once at init β€” avoids repeated lookups
        self.pad_id = self.tokenizer.token_to_id("<pad>")
        self.unk_id = self.tokenizer.token_to_id("<unk>")
        self.bos_id = self.tokenizer.token_to_id("<bos>")
        self.eos_id = self.tokenizer.token_to_id("<eos>")

    def encode(self, text: str) -> list[int]:
        """Encode a string to a list of integer token IDs."""
        return self.tokenizer.encode(text).ids

    def decode(self, ids: list[int]) -> str:
        """Decode a list of token IDs back to a string."""
        return self.tokenizer.decode(ids)

    def clean(self, text: str) -> str:
        """
        Remove byte-level BPE artefacts and normalise whitespace.
        """
        return text.replace("Ġ", " ").replace("Ċ", "\n").strip()

    @property
    def vocab_size(self) -> int:
        """Total vocabulary size (including special tokens)."""
        return self.tokenizer.get_vocab_size()