File size: 2,142 Bytes
a2b525b |
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 |
"""Compact Chess BPE Tokenizer: Train, Upload, Load, Inference"""
import os, json, rustbpe, tiktoken
from datasets import load_dataset
from huggingface_hub import HfApi, create_repo, upload_folder, hf_hub_download
REPO_ID = "ItsMaxNorm/bpess"
def train(vocab_size=4096, split="train[0:10000]"):
"""Train BPE tokenizer on chess moves."""
ds = load_dataset('angeluriot/chess_games', split=split)
tok = rustbpe.Tokenizer()
tok.train_from_iterator((' '.join(g['moves_custom']) for g in ds if g['moves_custom']), vocab_size)
return tok
def save(tok, path="./tokenizer"):
"""Save tokenizer files locally."""
os.makedirs(path, exist_ok=True)
ranks = tok.get_mergeable_ranks()
json.dump({bytes(k).decode('utf-8', errors='replace'): v for k, v in ranks},
open(f"{path}/vocab.json", 'w'), indent=2)
json.dump({"pattern": tok.get_pattern(), "vocab_size": tok.vocab_size},
open(f"{path}/config.json", 'w'))
return path
def upload(tok, repo_id=REPO_ID, private=False):
"""Upload tokenizer to HuggingFace Hub."""
path = save(tok)
try: create_repo(repo_id, private=private)
except: pass
HfApi().upload_folder(folder_path=path, repo_id=repo_id)
print(f"Uploaded: https://huggingface.co/{repo_id}")
def load_tiktoken(repo_id=REPO_ID):
"""Load tokenizer from HuggingFace as tiktoken Encoding."""
config = json.load(open(hf_hub_download(repo_id, "config.json")))
vocab = json.load(open(hf_hub_download(repo_id, "vocab.json")))
return tiktoken.Encoding(
name="chess", pat_str=config["pattern"],
mergeable_ranks={k.encode('utf-8', errors='replace'): v for k, v in vocab.items()},
special_tokens={}
)
if __name__ == "__main__":
# Train & Upload
tok = train(vocab_size=4096, split="train[0:10000]")
print(f"Trained: {tok.vocab_size} tokens")
upload(tok, REPO_ID)
# Load & Inference
enc = load_tiktoken(REPO_ID)
test = "w.♘g1♘f3.. b.♟c7♟c5.. w.♙d2♙d4.."
ids = enc.encode(test)
print(f"Encoded: {ids[:10]}... ({len(ids)} tokens)")
print(f"Decoded: {enc.decode(ids)}")
|