Нахождение пропущенных пробелов

Проект для восстановления пробелов в тексте без пробелов с использованием каскадной модели BERT.

Описание

Проект реализует алгоритм восстановления пробелов в текстах, где пробелы были удалены. Использует два подхода:

• BERT для нахождения пробелов - предсказывает позиции пробелов, реализация в модуле space_restorer.bert
• BERT для 'склеивания' нарочно созданных пустот - находит токены, нарочно разделенные пробелом, реализация в модуле space_restorer.bert_gap

Зависимости

В качестве утилиты управления зависимостями использовался uv. Для запуска кода рекомендуется использовать его:

pip install uv # или pipx install uv
uv venv --python 3.13 .venv
uv sync

Можно установить через pip:

pip install -r requirements.txt # не тестировал
pip install -e . # обязательно!

Можно также использовать docker:

docker build -t space-restorer .
docker volume create models # понадобится для доступа к чекпоинтам моделей в контейнере

Для вариантов запуска контейнера см. Использование.

Использование

По умолчанию скрипты ожидают веса моделей в формате чекпоинтов (ссылка на скачивания через Яндекс.Диск или с использованием скрипта download.sh) по пути:

avito-ds-autumn/
└── models/
    ├── checkpoint-gap       # Для модели склейки
    └── checkpoint-spaces    # Для модели восстановления пробелов

Note

Перечисленные ниже команды можно запускать внутри docker контейнера. Для этого необходимо запустить контейнер в интерактивном режиме и установить веса моделей, если это первый запуск:

docker volume create models # если еще не сделано
docker run -it --name space-restorer-container -v models:/app/models space-restorer /bin/bash
./download.sh # при первом запуске

Чтобы посмотреть формат ожидаемых данных для скриптов (модулей):

# Преобразовать входной файл (как на stepik) в выходной файл
uv run -m space_restorer --help # python -m space_restorer --help

# Восстановление пробелов для одного текста
uv run -m space_restorer.cascade --help # python -m space_restorer.cascade --help

Пример использования:

# Для файла
uv run -m space_restorer --max-tries 3 --save-path data/output/output.txt data/input/dataset.txt

# Для текста
uv run -m space_restorer --max-tries 3 "Приветмир"

Warning

В файле dataset.txt ожидается, что все строки, использующие запятую (,), обернуты в кавычки ("). Датасеты на stepik находятся в 'некорректном' формате.

Для включения недетерминированного режима:

uv run -m space_restorer --max-tries 5 --min_tries 2 --spaces 0.2 --save-path data/output/output.txt data/input/dataset.txt

REST API сервис

Доступно также использование REST API сервиса (используется FastAPI). Для этого необходимо запустить docker контейнер в обычном режиме:

docker volume create models # если еще не сделано
docker run -p 8000:8000 --name space-restorer-container -v models:/app/models space-restorer # веса скачаются автоматически

Пример использования сервиса (имеет единственный endpoint):

curl -X POST "http://localhost:8000/restore" -H "Content-Type: application/json" -d '{"text": "Примертекстаспропущеннымипробелами"}'

Валидация

Валидация по ключевой метрике F1 score производилась для разных значений гиперпараметров max_tries, min_tries, spaces. Ниже приведены графики на выборке обьемом 100 текстов (~2200 предложений, датасет ru_news). Для каждого параметра производилось усреднение по всем значениям остальных параметров.

Оптимальные значения параметров:

Время	spaces	max_tries	min_tries	f1
Быстро	0.0	3	0	~0.875
Любое	0.2	15	7	~0.89

Подход к решению

Предложенный алгоритм к поставленной задаче можно описать следующим образом:

1. Исходная последовательность символов приводится к виду, где каждое слово - одна буква (пример: "П р и в е т м и р")
2. Модель, обученная предсказывать токены текста, которые не должны быть разделены пробелом - склейка (пример: "П р и в е т м о й м и р" -> "Привет мой мир")
3. Поскольку 'склеивание' может произойти неидеально, по тексту проходит вторая модель, обученная находить токены, перед которыми должен стоять пробел (пример: "Привет моймир" -> "Привет мой мир")
4. Шаги 2 и 3 можно повторить несколько раз для большей надежности (можно не повторять)

• Между повторениями можно использовать недетерминированную эвристику - разорвать часть слов пробелами (случайно), чтобы затем повторно их склеить (об этом - далее)

Почему такой подход

В рамках решения поставленной задачи было выявлено следующее - с помощью классической bert модели (используя ее как TokenClassificator) нельзя полностью решить поставленную задачу - в исходной последовательности (подряд идущие символы) токенизация может произойти так, что токен находится на границе двух слов. В таком случае модель никоим образом не может верно разделить слова.

Модель 'склейки' лишена указанной выше проблемы, однако обладает средней точностю и высокой полнотой (часто склеивает лишние токены). Поэтому было решено использовать каскадный подход в своем решении.

Недетерминированное разбиение слов между циклами применения двух моделей призвано 'разбить' токены, находящие на границе двух слов. Если разбиение излишнее, 'полная' модель склейки устранит это.

Еще пару слов

Модель нахождения токенов, перед которыми находится пробел, независима от токенизатора и может использовать, например, DeepPavlov RuBERT. Однако токенизатор обычной мультиязыковой модели (SequencePiece) показал себя лучше.

Модель 'склейки' полностью зависит от токенизатора SequencePiece.

Системные требования

Модель обучена на системе со следующими характеристиками:

• GPU: NVIDIA 5060 Ti 16GB
• CPU: AMD Ryzen 9700x
• RAM: 32GB DDR5

По умолчанию запускается на GPU, можно принудительно запустить на CPU с указанием флага --cpu. Для запуска на GPU необходимо минимум 4GB VRAM. Для более быстрой работы алгоритма можно уменьшить максимальное число каскадных повторений с помощью флага --max-tries.

Упоминания

Проект использует предобученные модели BERT (мультиязычная), адаптированные для задачи восстановления пробелов. Предобученные веса из библиотеки huggingface transformers, название базовой модели - xlm-roberta-base.

Модели обучались на открытых датасетах IlyaGusev/picabu и IlyaGusev/ru_news.