WIP on refactoring and scripts reorganizing

.gitignore

@@ -1,3 +1,6 @@
+# macOS stuff
+.DS_Store
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -127,3 +130,7 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# Project
+notebooks
+data

docs/architecture.md

@@ -1,5 +1,39 @@
 # Архитектура модели YOGURT
-...
+
+## Вариант 1
+Есть набор предложений, в каждом предложении есть по меньшей мере одно слово, в котором один специальный символ нужно заменить на другой.
+Символов в одном слове может быть несколько, слов для замены в предложении тоже. Проблема в том, что менять символ нужно в зависимости от контекста, т. е. от окружающих слов. Правил очень много, формализовать все практически нереально, поэтому смотрю сразу в сторону DL и обучения с учителем на парах предложений до и после замены.
+
+Пока предполагаю, что будут две модели:
+ 1. Энкодер (уровня слов): кодирует предложение, чтобы отразить в нем связи между словами, семантический смысл и т. п.
+ 2. Декодер (уровня символов): получает на вход вектор-контекст слов от энкодера и закодированное слово для замены*
+
+Само слово для декодера решил *кодировать вот так:
+abcdef (допустим меняем f на z) -> [2, 0, 0, 0, 0, 0, 1], где 2 — позиция слова в тексте, 1 — маркер символа, 0 — паддинг.
+На выход от декодера в таком случае жду вектор [0, 0, 0, 0, 0, 0.75], где 0.75 — вероятность (условно) замены f на z.
+Поверх этого бинарная кросс-энтропия и классические метрики классификации. Как я это вижу: «менять ли символ f на z».
+
+Для первого приближения к решению пока могу не рассматривать слова, где символов несколько (в большинстве он все же один).
+Не исключаю, что я вообще лажу придумал и перемудрил от незнания, поэтому жду, что опытные товарищи укажут на это…
+
+## Вариант 2
+Если все правки, которые нужно внести в текст, сводятся к замене символа на символ (без вставки и удаления), то кажется, что seq2seq для этой задачи overshoot, и с ней вполне сможет справиться более простая модель формата sequence tagging. Конкретно по архитектуре это может быть и cnn, и rnn, и трансформер, смотря насколько сильна и сложна зависимость от контекста.
+Подавал в неё я бы тупо последовательность символов (всё предложение), а на выходе каждого токена требовал бы предсказать распределение над теми символами, которые там на самом деле должны быть.
+
+--
+Вот к чему-то такому склоняюсь. Смущает только, что по факту из всего предложения может быть 1-2 слова для замены, тогда модели придётся на каждый токен выдавать нулевые векторы. Но согласен, по крайней мере это проще и понятнее, чем мои эвристики по кодированию входа. Спасибо!
+
+--
+
+Ну и пусть выдает, жалко что ли) ведь ей же все равно приходится принимать решение, надо ли заменять каждый из этих токенов или нет, так пусть принимает его в явном виде.
+--
+Кстати правильно же понимаю, что в вашем предложении нельзя готовые эмбеддинги для слов поиспользовать, модель должна сама изучить зависимости между цепочками символов?
+--
+Если очень хочется, можно в этой модели объединить представления слов и символов: например, на первом слое конкатенировать эмбеддинг символа (обучаемый) с эмбеддингом слова (предобученным).
+
+Но вообще, я не уверен, что при правке текста на уровне символов эмбеддинги слов вообще полезны: во-первых, эмбеддинги слов обычно ничего не знают о написании этих слов (если это не fasttext), а во вторых, при замене символов (чем вы собираетесь заниматься), слово становится другим - и, скорее всего, оно будет OOV для предобученных эмбеддингов.
+
+Впрочем, я могу ошибаться. Можете привести примеры текстов из вашей задачи?)
 
 ## Полезные ссылки
 * [StressRNN](https://github.com/Desklop/StressRNN)

docs/data.md

@@ -0,0 +1,14 @@
+# Данные для обучения
+
+Модель обучена на данных из русской Википедии. Скрипты для получения данных находятся в директории `scripts`. Текущий датасет собран следующей командой:
+```shell
+python scripts/extract_sentences_from_wiki.py -j 8 -s 50 -n 1000000
+```
+Затем произведена первичная очистка текста:
+```shell
+python scripts/preprocess_sentences.py -j 4
+```
+Финальный датасет собирается следующей командой:
+```shell
+python scripts/compile_dataset.py
+```

readme-ru.md

@@ -1,7 +1,7 @@
 # Йогурт — умное восстановление буквы «ё» в русских текстах
 
 ## Проблема
-При работе с текстом на русском языке часто возникает потребность восстановить в словах оригинальную букву «ё», например, для улучшения читаемости или единообразия стиля письма. Кроме того, крайне желательно учитывать разницу между «е» и «ё» при обработке текстов в задачах машинного обучения, в частности, при распознавании и синтезе речи. 
+При работе с текстом на русском языке часто возникает потребность восстановить в словах оригинальную букву «ё», например, для улучшения читаемости или единообразия стиля письма. Кроме того, крайне желательно учитывать разницу между «е» и «ё» при обработке текстов в задачах машинного обучения, в частности, при распознавании и синтезе речи.
 
 Однако в открытых инструментах для восстановления буквы «ё», таких как [`eyo-kernel`](https://github.com/e2yo/eyo-kernel), при замене слов не учитывается их контекст. Это приводит к тому, что слова можно восстановить только в однозначных случаях, когда слово без «ё» не употребляется, например:
 ```

requirements.txt

@@ -0,0 +1,5 @@
+corus==0.9.0
+numpy==1.23.3
+pandas==1.5.0
+razdel==0.5.0
+tqdm==4.64.1

scripts/download_russian_wiki_dump.sh

@@ -2,6 +2,5 @@
 
 set -e
 
-cd ../;
 mkdir -p data;
 curl -Lo data/ruwiki-latest-pages-articles.xml.bz2 https://dumps.wikimedia.org/ruwiki/latest/ruwiki-latest-pages-articles.xml.bz2;

scripts/extract_sentences_from_wiki.py

@@ -0,0 +1,102 @@
+"""A script that extracts the sentences with `Ё` letter from the Russian Wikipedia dump."""
+
+import argparse
+import logging
+import multiprocessing as mp
+
+from typing import List
+
+from corus.sources.wiki import WikiRecord, load_wiki
+from tqdm import tqdm
+
+from src import utils
+
+
+# suppress warnings from Wiki extractor
+logging.getLogger().setLevel(logging.ERROR)
+
+
+def job(records: List[WikiRecord]) -> List[str]:
+    sentences = []
+    for record in records:
+        normalized = utils.normalize_wiki_text(record.text)
+        sentences.extend(utils.extract_unique_yo_segments(normalized, repl=' '))
+    return sentences
+
+
+def aggregate_job_results(pool: mp.Pool, jobs: List[List[WikiRecord]]) -> List[str]:
+    results = pool.imap_unordered(job, jobs)
+    return sum(results, [])
+
+
+def main(args: argparse.Namespace):
+    assert args.num_sentences is None or args.num_sentences > 0
+    wiki = load_wiki(args.wiki_path)
+
+    sentences = []
+    with mp.Pool(args.njobs) as pool, tqdm(
+        total=args.num_sentences,
+        leave=True,
+        desc='Extracting `Ё` sentences from wiki records',
+        dynamic_ncols=True,
+    ) as progress:
+        jobs = []
+        for records in utils.batch(wiki, args.jobsize):
+            if len(jobs) < args.njobs:
+                jobs.append(records)
+            else:
+                found = aggregate_job_results(pool, jobs)
+                progress.update(len(found))
+                sentences.extend(found)
+                jobs.clear()
+
+            if args.num_sentences and len(sentences) >= args.num_sentences:
+                sentences = sentences[:args.num_sentences]
+                progress.update()
+                progress.close()
+                break
+
+    with open(args.save_path, 'w', encoding='utf-8') as file:
+        for sentence in sentences:
+            file.write(sentence + '\n')
+
+    print(f'File saved to: {args.save_path}')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        '-w', '--wiki-path',
+        help='a path to wiki dump',
+        default='data/ruwiki-latest-pages-articles.xml.bz2'
+    )
+    parser.add_argument(
+        '-f', '--save-path',
+        help='a filepath to save the sentences',
+        default='data/ruwiki-yo-sentences.txt'
+    )
+    parser.add_argument(
+        '-j', '--njobs',
+        metavar='INT',
+        type=int,
+        default=4,
+        help='a number of parallel jobs',
+    )
+    parser.add_argument(
+        '-s', '--jobsize',
+        metavar='INT',
+        type=int,
+        default=10,
+        help='a number of documents for a single job',
+    )
+    parser.add_argument(
+        '-n', '--num-sentences',
+        metavar='INT',
+        type=int,
+        default=None,
+        help='a hard limit of sentences to gather'
+    )
+    main(parser.parse_args())

scripts/preprocess_sentences.py

@@ -0,0 +1,59 @@
+"""A script for initial preprocessing/cleaning of `Ё` sentences."""
+
+import argparse
+import os
+import multiprocessing as mp
+
+from tqdm import tqdm
+from src import utils
+
+
+def print_filesize(filepath: str):
+    print(f'File size is {os.stat(filepath).st_size / (1024 ** 3):.2f} GB')
+
+
+def main(args: argparse.Namespace):
+    with open(args.data_path) as file:
+        data = file.readlines()
+
+    print('Opened initial data.')
+    print_filesize(args.data_path)
+
+    with mp.Pool(8) as pool, tqdm(
+        pool.imap_unordered(utils.normalize_wiki_text, data),
+        total=len(data),
+        desc='Cleaning `Ё` sentences'
+    ) as progress:
+        new_data = list(progress)
+
+    with open(args.save_path, 'w') as file:
+        for sentence in filter(bool, new_data):
+            file.write(sentence + '\n')
+
+    print('Saved results to a new file.')
+    print_filesize(args.save_path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        '-d', '--data-path',
+        help='a path to a raw wiki sentences TXT file',
+        default='data/ruwiki-yo-sentences.txt'
+    )
+    parser.add_argument(
+        '-f', '--save-path',
+        help='a filepath to save the cleaned sentences',
+        default='data/ruwiki-yo-sentences-preprocessed.txt'
+    )
+    parser.add_argument(
+        '-j', '--njobs',
+        metavar='INT',
+        type=int,
+        default=4,
+        help='a number of parallel jobs',
+    )
+    main(parser.parse_args())