from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.processors import TemplateProcessing

# Initialize a tokenizer
tokenizer = Tokenizer(BPE(unk_token="[UNK]"))

# Use the byte level pre-tokenizer
tokenizer.pre_tokenizer = ByteLevel()

# Customize training with a BpeTrainer
trainer = BpeTrainer(special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"])

# Path to the file(s) for training the tokenizer
files = ["/home/saicharan/Downloads/chembl.csv"]

# Train the tokenizer
tokenizer.train(files, trainer)

# Optionally, you can customize the post-processing to add special tokens
tokenizer.post_processor = TemplateProcessing(
    single="[CLS] $A [SEP]",
    pair="[CLS] $A [SEP] $B:1 [SEP]:1",
    special_tokens=[
        ("[CLS]", tokenizer.token_to_id("[CLS]")),
        ("[SEP]", tokenizer.token_to_id("[SEP]")),
    ],
)

# Save the tokenizer
tokenizer.save("/home/saicharan/Downloads/chembl_bpe_tokenizer.json")