WordPiece

Репозиторий содержит два самых быстрых алгоритма токенизации WordPiece.

1. Linear. Оптимальная асимптотика по времени -- $O(n)$, память -- $O(n)$ (с константой около 12-20), реализация эффективнее, чем в статье "Fast WordPiece" (2021, torch). Многопоточная реализация.

2. Fast. Асимптотика -- $O(nm)$, где $m$ -- максимальная длина слова в словаре. Многопоточная реализация, на практике эффективнее, чем другие библиотеки минимум в 5 раз. Обгоняет Linear в 1.5-10 раз.

Check out our benchmark results.

Linear Algorithm

Возьмем строку s#t1#t2#...tk, построим на ней суфмас. Решетки обязательно надо вставить (символ, который не встречается).

Наша задача для каждого суффикса s найти самую длинную ti которая равна префиксу этого суффикса в суфмасе есть k интересных мест, там где начинаются ti, у каждой есть длина |ti| хотим для всех позиций суфмаса найти самую длинную хорошую интересную позицию, такую что лцп до нее хотя бы длина этой позиции (такая интересная позиция будет называться хорошей, то есть та, до которой лцп хотя бы длина) для этого найдем самую длинную интересную хорошую позицию слева и справа для каждой позиции в суфмасе независимо, дальше возьмем бОльшую из двух ищем слева, справа аналогично. пусть есть две интересных позиции в суфмасе с индексами i<j, причем |i| >= |j|. В предположении что все ti различны, можно утверждать, что i уже никогда не будет хорошей для всех позиций суфмаса >j, поскольку лцп (i;j) строго меньше |i| (если |i| > |j|, то она просто не больше |j|, а если равны, то так как строки различные, то строго меньше). Это означает, что если мы фиксируем позицию суфмаса e, то среди интересных позиций слева от e нас интересуют только рекорды, то есть только убывающий стэк. алгоритм: идем сканлайном слева направо, храним стэк хороших интересных позиций i1, i2, .. in, причем |i1| < |i2| < ... < |in|, изначально стэк пустой. пришли в очередную позицию e, пусть lcp(e-1, e) = x. Тогда надо выкинуть все элементы с конца стэка, что |in| > x (очевидно они стали плохими позициями). Теперь если e сама является хорошей позицией, то надо ее добавить в стэк. Очевидно что lcp(e-1, e) <= |e|, так что e будет максимумом в стэке, и инвариант возрастания сохраняется. в очередной позиции e самая длинная интересная хорошая справа позиция это in (потому что в стэке все хорошие позиции и только они).

Fast Algorithm

Стоя на позиции i возьмем подстроку [i, i + m), где m -- длина максимального слова в словаре. Проверим ее наличие в словаре-хешмапе. Если нашлось совпадение, то сохраним токен в ответ и сдвинем позицию. Если совпадение не нашлось, то уберем последний символ из подстроки. Повторяем пока подстрока не пуста. Если повторы дошли до пустой подстроки, то добавим UNK в ответ и сдвинем позицию до начала следующего слова.

Roadmap

1. интеграция в youtokentome;
2. версия linear для GPU - с помощью https://github.com/IlyaGrebnov/libcubwt
3. возможные оптимизации fast:
   • для маленьких алфавитов max_len можно брать по первой букве токена а не во всем словаре
   • сейчас бывают отрезки где хэш пересчитывается, хотя можно посчитать хэш и выбрасывать префикс (http://e-maxx.ru/algo/string_hashes#3)

Build

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release && make -C build

To build with OpenMP: -DCMAKE_USE_OPENMP=On, Sanitizers: -DCMAKE_USE_SANITIZERS;

Prepare benchmark

apt install wget bzip2 perl cmake make
mkdir -p data
wget -O data/vocab.txt https://huggingface.co/bert-base-cased/resolve/main/vocab.txt
python3 -m venv venv && source venv/bin/activate
pip3 install -r tests/requirements.txt
cmake -S . -B build -DCMAKE_BUILD_TYPE=Release

For ARM arch use tests/requirements-arm.txt (no tensorflow).

Run benchmark

source venv/bin/activate && make -C build && python3 tests/speed_test.py --langs en ru ja zh --vocab data/vocab.txt --corpus_size 10 --n_threads 8