from datasets import load_dataset
from tokenizers import (
    decoders,
    models,
    normalizers,
    pre_tokenizers,
    processors,
    trainers,
    Tokenizer,
    Regex,
)
from transformers import PreTrainedTokenizerFast, PreTrainedTokenizerBase
from tqdm import tqdm

dataset = load_dataset(
    "parquet", data_dir="Mxode/IndustryCorpus-Subset-zh-en", split="train")
dataset = dataset.shuffle(seed=3407)

ds = dataset[:1000000]
ds_val = dataset[-10000:]
char_len = sum(len(x) for x in ds_val['text'])


def get_training_corpus():
    for i in range(0, len(ds), 1000):
        yield ds["text"][i: i + 1000]


def train():
    tokenizer = Tokenizer(models.BPE())
    tokenizer.normalizer = normalizers.NFC()
    tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
        pre_tokenizers.Split(
            pattern=Regex(
                "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"),
            behavior="isolated",
            invert=False,
        ),
        pre_tokenizers.ByteLevel(
            add_prefix_space=False,
            use_regex=False,
            trim_offsets=False
        )
    ])
    trainer = trainers.BpeTrainer(
        vocab_size=16000,
        special_tokens=["<|endoftext|>", "<|im_start|>", "<|im_end|>"]
    )
    tokenizer.train_from_iterator(get_training_corpus(), trainer=trainer)
    tokenizer.post_processor = processors.ByteLevel(
        add_prefix_space=False,
        use_regex=False,
        trim_offsets=False
    )
    tokenizer.decoder = decoders.ByteLevel(
        add_prefix_space=False,
        use_regex=False,
        trim_offsets=False
    )
    wrapped_tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        bos_token="<|endoftext|>",
        eos_token="<|im_end|>",
        pad_token="<|endoftext|>",
        model_max_length=4096,
        clean_up_tokenization_spaces=False,
        errors="replace",
        split_special_tokens=False,
    )
    wrapped_tokenizer.chat_template = """{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"""
    wrapped_tokenizer.save_pretrained(
        'Mxode/Bilingual-Tokenizer/BilingualTokenizer-16K')
    return wrapped_tokenizer


def eval(tokenizer: PreTrainedTokenizerBase):
    def get_compress_len(tokenizer):
        return sum(len(tokenizer(x, return_tensors=None)['input_ids']) for x in tqdm(ds_val['text']))

    compress_len = get_compress_len(tokenizer)
    compression_rate = compress_len / char_len * 100
    print(f'{len(tokenizer):<40} {compression_rate:.2f}%')


if __name__ == "__main__":
    tokenizer = train()
    eval(tokenizer)