init bpess.py

bpess.py

@@ -0,0 +1,54 @@
+"""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)}")