Upload train.py

train.py

@@ -0,0 +1,87 @@
+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)