from pathlib import Path

import tiktoken
from transformers.integrations.tiktoken import convert_tiktoken_to_fast
from transformers import PreTrainedTokenizerFast

out_dir = Path("hf_tokenizer")
out_dir.mkdir(exist_ok=True)

# Convert
encoding = tiktoken.get_encoding("cl100k_base")
convert_tiktoken_to_fast(encoding, str(out_dir))

# Load the generated tokenizer
tokenizer = PreTrainedTokenizerFast(
    tokenizer_file=str(out_dir / "tokenizer.json"),
    bos_token="<|endoftext|>",
    eos_token="<|endoftext|>",
    unk_token="<|endoftext|>",
)

# Save all Hugging Face tokenizer files
tokenizer.save_pretrained(out_dir)

print(f"Saved tokenizer to {out_dir}")