WikiTextデータセットを手元に置いておきたい

この記事について

LLMの勉強のために実際にモデルの性能を測定して検証してみようと思っています。評価用のデータセットは色々有りますが、出来れば手元に保存しておけるぐらいのサイズで論文や様々な検証でしばしば取り上げられるデータセットをと考えた時にwikitext2データセットを利用することにしました。Wikipediaから抽出した100M(1億)トークン規模のテキストデータセットです。その保存方法ですが、僅か数行のコードで保存できます。

従来は直接ダウンロードも出来ていた

以前はSalesforceがAWS S3上でデータセットのオリジナルのファイルを公開していました。ところが、最近公開を取りやめてしまったようで、ダウンロード出来なくなりました。

直接ダウンロードできなくなったが。。。

直接ファイルをダウンロードできなくなりましたが、Pythonのdatasetsライブラリからはwikitext2データセットをロードして利用出来ます。これを上手く応用することでデータセットを保存できます。

実際にやってみる

Pythonの簡単なスクリプトを作成してデータセットを保存してみます。GitHub Copilot君にスクリプトの実装案を提案させてみます。その前に環境構築です。筆者はuvを用いて構築しています。

mkdir wikitext2
cd wikitext2
uv venv --python 3.14
uv pip install datasets

Pythonスクリプトのコードは以下です。

from pathlib import Path

import datasets


def save_wikitext2():
    dataset = datasets.load_dataset("wikitext", "wikitext-2-raw-v1")
    save_path = Path("wikitext-2-raw-v1")
    save_path.mkdir(parents=True, exist_ok=True)
    for split in ("train", "validation", "test"):
        with open(save_path / f"{split}.txt", "w", encoding="utf-8") as f:
            for line in dataset[split]["text"]:
                f.write(line)


if __name__ == "__main__":
    save_wikitext2()

実はこのままではエラー終了します。エラーメッセージを抜粋引用すると。。。

'hf://datasets/wikitext@b08601e04326c79dfdd32d625aee71d232d685c3/.huggingface.yaml'. Repository id must be 'namespace/name', got 'wikitext'.

namespace/nameという表記が必要なので、Salesforce/wikitextという表記が必要です。

from pathlib import Path

import datasets


def save_wikitext2():
    dataset = datasets.load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1")
    save_path = Path("wikitext-2-raw-v1")
    save_path.mkdir(parents=True, exist_ok=True)
    for split in ("train", "validation", "test"):
        with open(save_path / f"{split}.txt", "w", encoding="utf-8") as f:
            for line in dataset[split]["text"]:
                f.write(line)


if __name__ == "__main__":
    save_wikitext2()

これでwikitext2データセットをtxtファイルで保存できます。train.txt, test.txt, validation.txtという3種類のtxtファイルとして保存されます。

まとめ

wikitext2データセットを無事手元に保存できました。スクリプトだけGitHubに上げておけば、いつでも再保存出来ます。