import json, os, torch
class ChessTokenizer:
    def __init__(self, vocab):
        self.vocab = vocab
        self.id_to_token = {v: k for k, v in vocab.items()}
        self.pad_token_id = vocab.get("[PAD]", 0)
        self.bos_token_id = vocab.get("[BOS]", 1)
        self.eos_token_id = vocab.get("[EOS]", 2)
        self.unk_token_id = vocab.get("[UNK]", 3)
    @property
    def vocab_size(self): return len(self.vocab)
    def _convert_token_to_id(self, token): return self.vocab.get(token, self.unk_token_id)
    def pad(self, encoded_inputs, padding=True, max_length=None, pad_to_multiple_of=None, return_tensors=None, **kwargs):
        max_len = max(len(x["input_ids"]) for x in encoded_inputs)
        padded = []
        for x in encoded_inputs:
            ids = x["input_ids"] + [self.pad_token_id] * (max_len - len(x["input_ids"]))
            padded.append(ids)
        if return_tensors == "pt": return {"input_ids": torch.tensor(padded, dtype=torch.long)}
        return {"input_ids": padded}
    def save_pretrained(self, path):
        os.makedirs(path, exist_ok=True)
        with open(os.path.join(path, "vocab.json"), "w") as f: json.dump(self.vocab, f, indent=2)
    @classmethod
    def from_pretrained(cls, path):
        with open(os.path.join(path, "vocab.json"), "r") as f: return cls(json.load(f))