Spaces:
Sleeping
Sleeping
File size: 2,422 Bytes
a5f2e46 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | from tokenizers import Tokenizer
from tokenizers import models, pre_tokenizers, trainers, processors, normalizers
from .move_tokenizer import MoveTokenizer
from .position_tokenizer import PostionTokenizer
def create_bpe_tokenizer(
vocab_size: int = 2000,
) -> tuple[Tokenizer, trainers.BpeTrainer]:
"""
Create a BPE tokenizer for chess moves.
The tokenizer uses Byte-Pair Encoding (BPE) with ByteLevel pre-tokenization
to handle the unique structure of chess moves. The <STEP> token is used to
explicitly separate each move, ensuring clear move boundaries.
Special tokens included:
- <PAD>: Padding token
- <START>: Start of sequence token
- <END>: End of sequence token
- <STEP>: Move separator token
- <1-0>: White wins
- <0-1>: Black wins
- <1/2-1/2>: Draw
- <UNK>: Unknown token
Args:
vocab_size: Size of the vocabulary (default: 2000)
Returns:
A Tokenizer object and trainer
"""
# Initialize a tokenizer with BPE model
encoder = Tokenizer(models.BPE(unk_token="<UNK>"))
encoder.normalizer = normalizers.Sequence([]) # No normalization needed for chess moves
# remove the white space pre-tokenizer
encoder.pre_tokenizer = pre_tokenizers.Whitespace()
# Setup trainer with special tokens, including space
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
special_tokens=[
"<PAD>",
"<START>",
"<END>",
"<1-0>",
"<0-1>",
"<1/2-1/2>",
"<UNK>",
"<STEP>",
],
show_progress=True,
min_frequency=2,
initial_alphabet=[
"a",
"b",
"c",
"d",
"e",
"f",
"g",
"h",
"1",
"2",
"3",
"4",
"5",
"6",
"7",
"8",
"q",
"k",
"r",
"b",
"n", # Piece notations
],
)
# Add post-processing to add special tokens
encoder.post_processor = processors.TemplateProcessing(
single="<START> $A <END>",
special_tokens=[
("<START>", 1),
("<END>", 2),
],
)
return encoder, trainer
__all__ = ["create_bpe_tokenizer", "MoveTokenizer", "PostionTokenizer"]
|