import json
from pathlib import Path
from transformers import PreTrainedTokenizerFast


class iceAudioTokenizer(PreTrainedTokenizerFast):
    model_input_names = ["input_ids", "attention_mask"]
    vocab_files_names = {"tokenizer_file": "tokenizer.json"}

    def __init__(self, tokenizer_file=None, **kwargs):
        super().__init__(tokenizer_file=tokenizer_file, **kwargs)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *init_inputs, **kwargs):
        path = Path(pretrained_model_name_or_path)
        if path.is_dir() and (path / "tokenizer.json").is_file():
            cfg_path = path / "tokenizer_config.json"
            cfg = json.loads(cfg_path.read_text(encoding="utf-8")) if cfg_path.is_file() else {}
            allowed = {"bos_token","eos_token","unk_token","pad_token","sep_token","cls_token","mask_token",
                       "extra_special_tokens","additional_special_tokens","chat_template","model_max_length",
                       "padding_side","truncation_side","clean_up_tokenization_spaces","split_special_tokens",
                       "add_prefix_space","errors"}
            init_kwargs = {k:v for k,v in cfg.items() if k in allowed}
            init_kwargs.update(kwargs)
            init_kwargs.pop("tokenizer_file", None)
            return cls(tokenizer_file=str(path / "tokenizer.json"), **init_kwargs)
        return super().from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)


__all__ = ["iceAudioTokenizer"]