from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from datasets import load_dataset
import os

def download_to_txt(dataset_name, config_name, split, text_field, out_file, limit):
    print(f"⬇️ Скачиваем {out_file}...")
    # Загружаем стримингом, но сразу пишем на диск
    ds = load_dataset(dataset_name, config_name, split=split, streaming=True)
    
    with open(out_file, "w", encoding="utf-8") as f:
        for i, item in enumerate(ds):
            if i >= limit: break
            text = item.get(text_field, "")
            if text.strip():
                f.write(text + "\n")
    print(f"✅ Готово: {out_file}")


download_to_txt("wikimedia/wikipedia", "20231101.en", "train", "text", "en.txt", 30000)

download_to_txt("wikimedia/wikipedia", "20231101.ru", "train", "text", "ru.txt", 20000)

download_to_txt("wikimedia/wikipedia", "20231101.zh", "train", "text", "zh.txt", 10000)

download_to_txt("codeparrot/codeparrot-clean", None, "train", "content", "code.txt", 5000)


special_tokens = [
    "<unk>", "<s>", "</s>", "<pad>", 
    "<|im_start|>", "<|im_end|>", 
    "<|think|>", "<|/think|>"
]
tokenizer = Tokenizer(BPE(unk_token="<unk>"))
tokenizer.pre_tokenizer = ByteLevel()
tokenizer.decoder = ByteLevelDecoder()
trainer = BpeTrainer(vocab_size=64000, special_tokens=special_tokens, min_frequency=2)


files = ["en.txt", "ru.txt", "zh.txt", "code.txt"]
tokenizer.train(files, trainer)


tokenizer.save("tokenizer.json")

# Удаляем временные файлы, чтобы не забивать диск Colab
for f in files:
    if os.path.exists(f):
        os.remove(f)
print("🧹 Временные файлы очищены.")


print("\n--- ТЕСТ ТОКЕНИЗАТОРА ---")
tok = Tokenizer.from_file("tokenizer.json")
test_texts = {
    "Английский": "The quick brown fox jumps over the lazy dog.",
    "Русский": "Привет, как дела? Это наш TinyCyclic.",
    "Китайский": "这是一个测试，看看我们的分词器是否理解中文。",
    "Код": "def hello_world(): print('Hello')",
    "Спецтокены": "<|im_start|>assistant<|think|>Думаю...<|/think|>Ответ<|im_end|>"
}
for lang, text in test_texts.items():
    encoded = tok.encode(text)
    print(f"[{lang}] Токенов: {len(encoded.tokens)} -> {encoded.tokens}")