restoring-ancient-russian-texts/train_tokenizer.py at main · ukolchuga/restoring-ancient-russian-texts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import os

from tokenizers import BertWordPieceTokenizer
from transformers import BertTokenizerFast

CORPUS_FILE = "data/ancient_rus_ready_for_bert.txt"
VOCAB_SIZE = 50000
SAVE_DIR = "ancient_rus_tokenizer"


def main():
    if not os.path.exists(CORPUS_FILE):
        print(f"❌ Ошибка: Файл {CORPUS_FILE} не найден!")
        return

    print("⏳ Инициализация WordPiece токенизатора...")

    # clean_text=False -> мы уже сами всё идеально почистили
    # strip_accents=False -> ЗАПРЕЩАЕМ удалять титла!
    # lowercase=False -> мы уже сделали Sentence case
    tokenizer = BertWordPieceTokenizer(
        clean_text=False,
        handle_chinese_chars=False,
        strip_accents=False,
        lowercase=False,
    )

    special_tokens = [
        "[PAD]",
        "[UNK]",
        "[CLS]",
        "[SEP]",
        "[MASK]",
        "[GAP]",
        "[CTX_DAILY]",
        "[CTX_EPIC]",
        "[CTX_LIT]",
        "[CTX_LEGAL]",
        "[CTX_CHURCH]",
        "[CTX_SCIENCE]",
        "·",
        ":",  # Защита пунктуации
    ]

    print(f"🧠 Начинаем обучение токенизатора на файле {CORPUS_FILE}...")
    print(f"Размер словаря: {VOCAB_SIZE} токенов")

    tokenizer.train(
        files=[CORPUS_FILE],
        vocab_size=VOCAB_SIZE,
        min_frequency=2,
        show_progress=True,
        special_tokens=special_tokens,
        wordpieces_prefix="##",
    )

    os.makedirs(SAVE_DIR, exist_ok=True)
    tokenizer.save_model(SAVE_DIR)

    print("⏳ Настройка обертки Transformers...")
    # Явно указываем путь к vocab.txt, который сгенерировал tokenizers
    fast_tokenizer = BertTokenizerFast(
        vocab_file=f"{SAVE_DIR}/vocab.txt",
        strip_accents=False,
        lowercase=False,
        clean_text=False,  # Не даем HF портить нашу чистку
    )

    special_tokens_dict = {
        "additional_special_tokens": [
            "[CTX_DAILY]",
            "[CTX_EPIC]",
            "[CTX_LIT]",
            "[CTX_LEGAL]",
            "[CTX_CHURCH]",
            "[CTX_SCIENCE]",
            "[GAP]",
            "·",
            ":",
        ]
    }
    fast_tokenizer.add_special_tokens(special_tokens_dict)

    # Сохраняем HF конфиг
    fast_tokenizer.save_pretrained(SAVE_DIR)

    print("\n" + "=" * 50)
    print("✨ WORDPIECE ТОКЕНИЗАТОР УСПЕШНО ОБУЧЕН ✨")
    print(f"Папка с файлами: {SAVE_DIR}/")
    print("=" * 50)

    print("\n🔍 ТЕСТИРОВАНИЕ:")
    test_texts = [
        "[CTX_DAILY] Поклонъ · ѿ · бориса · [GAP] · настасии съ бг҃омъ.",
        "[CTX_CHURCH] Преподобноисповѣдникъ моляшеся непрестанно.",
    ]

    for i, text in enumerate(test_texts, 1):
        print(f"\n--- Тест {i} ---")
        print(f"Оригинал: {text}")
        tokens = fast_tokenizer.tokenize(text)
        print(f"Токены:   {tokens}")
        encoded_ids = fast_tokenizer.encode(text)
        print(f"Декодинг: {fast_tokenizer.decode(encoded_ids)}")


if __name__ == "__main__":
    main()