🐷

LLM コーパス構築用: deduplicate-text-datasets で Exact Dedup するメモ

2024/01/15に公開

RefinedWeb で Suffix Array での Excact Dedup で使われている手法

Deduplicating Training Data Makes Language Models Better
(通称 Lee et al. 2021)

の実装 https://github.com/google-research/deduplicate-text-datasets を試します!

日本語いける?

中身的には, 文字はすべてバイト列として扱い(UTF-8 文字にはしない)処理しています.

したがって日本語も特別に対応しなくとも使えます.

tokenizer するのもできます(デフォは GPT2).
GPT2 tokenizer だと日本語は基本 UTF-8 byte fallback(3 bytes)になってしまうため, riina あたりの日本語トークナイザー使うとよいでしょう.

試す

まずはコアをビルドしておきます.

$ cargo build

日本語 wiki40b ダウンロード

もはや誰も使っていないのと思うのだが, tensorflow_data を使っているのでしぶしぶ tensorflow 関連インストールしておきます...
(データセットダウンロードだけなので, 一通り動くの試したら huggingface datasets に移行も検討してみましょう: https://zenn.dev/syoyo/articles/a83d44f4d2da1e )

日本語は wiki40b/ja 指定なのですが, これだと wiki40b/ja.test みたいなファイルが生成されてしまうので, scripts/load_dataset.py スクリプト修正してデータセット名だけ wiki40b/ja になるようにしておきます.

あとは README sample にあるように,

$ python3 scripts/load_dataset.py --data_dir ~/tensorflow_datasets --save_dir data --name wiki40b --split test

著者環境ではダウンロード終了&データ変換後 tensorflow のメモリリーク関連かなんかエラーがでてしまいましたが, ダウンロードと変換自体はできているので無視して OK でした.

suffix array 構築

tmp フォルダ作っておかないとエラーになるので, 先に tmp フォルダ作っておいたのち,

$ python3 scripts/make_suffix_array.py data/wiki40b.test

でいけます!

頻度をチェック

$ python scripts/count_occurrences.py --suffix data/wiki40b.test --query "東京"
Number of times present: 11952

Voila~

重複をチェック!

$ cargo run self-similar --data-file data/wiki40b.test --length-threshold 100 --cache-dir tmp/cache --num-threads 8
...
Duplicates found: 269594
Total time taken: 2236ms

👌

あとはいろいろいじって exact dedup できるでしょう!

さらなる高みへ...

tokenize する

日本語だと 1/2 ~ 1/3 に圧縮が期待されます.

メモリ効率よくしたい

全部オンメモリで扱うようで, C4(300GB)の処理には 600GB のメモリが必要とあります.

日本語特化 tokenizer である程度圧縮できるとは思いますが, mmap なり out-of-core で 128GB くらいで処理できるようにしたいところですね.

ただそれでも suffix array は一つのファイルにするなどしているため 100 GB ~ 1 TB 扱うのがめんどいです.

100 GB over の場合は自前で 10 GB 単位くらいで処理するような suffix array のコードを書くのが良いでしょう

他ツールとの連携

https://github.com/ChenghaoMou/text-dedup

には exact-dedup 部分に deduplicate-text-datasets の Rust 部分が使われているので, text-dedup 使うのも検討してみましょう.

Discussion