WITデータセットの日本語部分を取り出す

Wikipedia-based Image Text (WIT)という、108の言語のWikipediaページから画像とテキストを抽出したデータセットがGoogleから公開されています。
https://github.com/google-research-datasets/wit
https://ai.googleblog.com/2021/09/announcing-wit-wikipedia-based-image.html
https://arxiv.org/abs/2103.01913

このデータセットは画像とテキストのマルチモーダルモデルを学習させることを念頭に置いたもので、108の言語の中には日本語も含まれています。論文のFig.15によると日本語のデータ数 (train data)は1,084,871（ユニーク画像数874,699枚）とのことで、日本語に限ってもかなり大規模なものです。とかく日本語のテキストデータは英語と比べると種類も規模も限られている現状ですので、このような大規模データは貴重です。
そこで、データをダウンロードして日本語部分のみを取り出して、データの中身がどのようなものか観察してみたいと思います。

## データダウンロード
データはtsv形式で10個のgzファイルに分けられています。ダウンロードリンクはこちら
https://github.com/google-research-datasets/wit/blob/main/DATA.md
にあります。ファイル一つにつき2.5GBほどです。

作業は全てGoogle Colaboratory上で行いました。元ファイルはgoogleのストレージにあるので、Colaboratoryのストレージへのダウンロードは高速で一つ30秒ほどで済みます。Colabノートブックは[こちら](https://github.com/jovyan70/wit-download)に置いてあります。

以下のようなコードで、ファイルをダウンロード・解凍し、日本語データのみを保存します。
ファイル容量が大きいので、ファイルは1行ずつ読み込んでいます。
各行の先頭のカラムが言語を表しており、カラムはタブ区切りなので、基本的には先頭が`ja\t`に一致する行を抜き出せばよいのですが、途中に改行が入る文章があるせいで1レコードが複数行に分かれていることがあり、それだけではうまくいきません。そこで、全ての言語をリストアップし、行の始めが`{lang}\t`というパターンの時は新しいレコード、そうでなければ前の行の続きと判断した上で、日本語のレコードのみを選別しています。下記コードの`lang`を変えれば他の言語も取り出せます。

```python
import os
import re
from tqdm import tqdm

# データ保存先　マウント済みのDriveのディレクトリを指定
save_dir = '/content/drive/MyDrive/wit_data'
# 対象言語　日本語
lang = 'ja'
target_pattern = re.compile(lang + '\t')

# 全言語
language_set = {
    'af', 'an', 'ar', 'arz', 'ast', 'az', 'azb', 'ba', 'bar', 'be', 'be-tarask',
    'bg', 'bn', 'br', 'bs', 'ca', 'ce', 'ceb', 'ckb', 'cs', 'cv', 'cy', 'da',
    'de', 'el', 'en', 'eo', 'es', 'et', 'eu', 'fa', 'fi', 'fil', 'fr', 'fy',
    'ga', 'gl', 'hi', 'hr', 'hsb', 'ht', 'hu', 'hy', 'ia', 'id', 'io', 'is',
    'it', 'iw', 'ja', 'jv', 'ka', 'kk', 'kn', 'ko', 'la', 'lah', 'lb', 'lmo',
    'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'nan', 'nds', 'ne',
    'nl', 'nn', 'no', 'nv', 'oc', 'pa', 'pl', 'pt', 'qu', 'ro', 'ru', 'sco',
    'si', 'sk', 'sl', 'sq', 'sr', 'sr-Latn', 'sv', 'sw', 'ta', 'te', 'tg', 'th',
    'tr', 'tt', 'uk', 'ur', 'uz', 'vec', 'vi', 'vo', 'war', 'xmf', 'yue',
    'zh', 'zh-TW'
    }
all_pattern = re.compile('|'.join([lang + '\t' for lang in language_set]))

length_list = []
for idx in range(10):  # ファイルは10個に分かれている
  data_index = str(idx).zfill(5)
  gz_file_name = "wit_v1.train.all-" + data_index + "-of-00010.tsv.gz"

  # download and unzip
  download_link = "https://storage.googleapis.com/gresearch/wit/" + gz_file_name
  !wget $download_link
  !gzip -d $gz_file_name

  # extract the target language
  file_name = gz_file_name.replace('.gz', '')
  data = []
  with open(file_name, 'r') as f:
    col_names = next(f)  # 1行目はカラム名
    tmp_line = ""
    for raw_line in tqdm(f):
      if all_pattern.match(raw_line):
        # A new record starts in this line.
        # Check if the **previous** line is the target language
        if target_pattern.match(tmp_line):
          data.append(tmp_line)
        tmp_line = raw_line
      else:
        # This line is continued from the previous line.
        tmp_line = tmp_line + '\n' + raw_line
    # Check if the last line is the target language
    if target_pattern.match(tmp_line):
      data.append(tmp_line)

  # save the extracted file
  save_path = os.path.join(save_dir, file_name.replace('.tsv', '_' + lang + '.tsv'))
  with open(save_path, 'w') as f:
    f.write(col_names)
    f.writelines(data)

  # Remove the original file
  !rm $file_name

  length_list.append(len(data))
  print("#############################################")
  print("File ", file_name)
  print("Length ", len(data))
  print("#############################################")

print("Total # of records: ", sum(length_list))
```

得られたファイルの行数は以下のようになりました。

| ファイル番号 | 行数 |
| ----------: | ---: |
| 00000 | 108,048 |
| 00001 | 108,375 |
| 00002 | 108,398 |
| 00003 | 107,600 |
| 00004 | 108,049 |
| 00005 | 108,703 |
| 00006 | 108,454 |
| 00007 | 107,986 |
| 00008 | 108,424 |
| 00009 | 108,544 |
| 合計  | 1,082,581 |

合計の件数が論文の情報と少しずれていますが、いずれにせよ約100万行の大規模なデータです。

## データ観察
保存したデータを再読み込みして中身を見てみましょう。
日本語のみにしたことでファイル容量が大幅に削減されているので、pandasで読み込むことができます。

```python
dimport pandas as pd

data_index = '00000'
file_name = "wit_v1.train.all-" + data_index + "-of-00010_" + lang + ".tsv"
file_path = os.path.join(save_dir, file_name)
df = pd.read_table(file_path)
df.shape  # -> (108048, 17)
```
試しに一つの行の中身を見てみると以下のようになっています。

```json-doc
language :  ja
page_url :  https://ja.wikipedia.org/wiki/%E3%83%97%E3%83%AD%E3%82%B0%E3%83%AC%E3%82%B9M1-3
image_url :  https://upload.wikimedia.org/wikipedia/commons/0/07/ISS_Zvezda_module.jpg
page_title :  プログレスM1-3
section_title :  nan
hierarchical_section_title :  プログレスM1-3
caption_reference_description :  STS-106から撮影
caption_attribution_description :  ISS Zvezda service module with docked Progress M1-3 spacecraft, taken in September 2000 by STS-106 (NASA)
caption_alt_text_description :  nan
mime_type :  image/jpeg
original_height :  1104
original_width :  1632
is_main_image :  True
attribution_passes_lang_id :  False
page_changed_recently :  False
context_page_description :  プログレスM1-3は2000年8月6日に打ち上げられたプログレス補給船。国際宇宙ステーションへの補給活動に使用され、～～中略～～ 海上に落下した。
context_section_description :  プログレスM1-3（Progress M1-3）は2000年8月6日に打ち上げられたプログレス補給船。国際宇宙ステーションへの補給活動に使用され、～～中略～～ 海上に落下した。
 ```

最初のカラムは言語、続いてwikipediaのページURL (`page_url`)、該当画像のURL (`image_url`) があります。画像自体はデータセットに含まれていないので、使用する際には自分で個別にダウンロードする必要があります。
`mime_type`, `original_height`, `original_width`は画像ファイルの種類と大きさを表します。
`is_main_image`, `attribution_passes_lang_id`, `page_changed_recently`の3つはメタデータですが、`is_main_image`が重要で、該当画像がページのメイン画像かどうかを区別するものです。
テキストが入ったカラムは複数ありますが、実際のwikipediaページとの対応関係は下の画像の通りです。

![progressM1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/613488/8bd7ba5e-e25b-d2b7-66ab-76e2a6db04d2.png)

- `page_title`はそのままページのタイトル
- `hierarchical_section_title`は`is_main_image=True`の場合は`page_title`と同じ
- `caption_reference_description`は画像の下に付いているキャプション。この例だと、画像の直接的なキャプションである`caption_reference_description`よりも`page_title`の方が画像の説明として相応しい。
- `caption_attribution_description`はメインページには表示されていない情報で、画像をクリックすると表示される。今回の例だと、キャプションとして一番ふさわしい内容ではあるが、英語で書かれている。
- `context_page_description`は記事の本文に該当し、長い文章で複数の改行やタブを含むことも多い。
- `context_section_description`は`is_main_image=True`の場合は`context_page_description`と同じく本文。今回の記事だと（Progress M1-3）などの括弧書きが省略されていないところのみが`context_page_description`と異なる。

以上で見たように画像に対してキャプションの文章一つだけが単純に結びついているのではなく、ページタイトルや記事本文など画像との関係性が様々な複数のテキスト項目があるというのがWITデータセットの特徴です。[Google AI blog](https://ai.googleblog.com/2021/09/announcing-wit-wikipedia-based-image.html#:~:text=Contextual%20information%3A%20Unlike%20typical%20multimodal%20datasets%2C%20which%20have%20only%20one%20caption%20per%20image%2C%20WIT%20includes%20many%20page%2Dlevel%20and%20section%2Dlevel%20contextual%20information.)によると、このような文脈情報 (contextual information)があるのがWITデータセットの一つの特長とのことですが、実際にデータを使う時のことを考えると、どの項目を使えばよいか難しいところでもあると思います。

いくつかデータを見た限りでは、`is_main_image`がTrueかFalseかで、画像と各文章項目の関係性の傾向が異なっているように感じました。上で見せたのはTrueの場合だったので、Falseのデータの一つを見てみましょう。

```json-doc
language :  ja
page_url :  https://ja.wikipedia.org/wiki/%E6%B0%B4%E7%B4%A0%E5%8E%9F%E5%AD%90%E3%81%AB%E3%81%8A%E3%81%91%E3%82%8B%E3%82%B7%E3%83%A5%E3%83%AC%E3%83%BC%E3%83%87%E3%82%A3%E3%83%B3%E3%82%AC%E3%83%BC%E6%96%B9%E7%A8%8B%E5%BC%8F%E3%81%AE%E8%A7%A3
image_url :  https://upload.wikimedia.org/wikipedia/commons/4/40/R-2s-orbital.png
page_title :  水素原子におけるシュレーディンガー方程式の解
section_title :  具体的な値
hierarchical_section_title :  水素原子におけるシュレーディンガー方程式の解 / 具体的な値
caption_reference_description :  nan
caption_attribution_description :  English: Radial distribution functions in 2s orbital 日本語: 2s軌道の動径関数
caption_alt_text_description :  nan
mime_type :  image/png
original_height :  960
original_width :  1280
is_main_image :  False
attribution_passes_lang_id :  True
page_changed_recently :  True
context_page_description :  本項、水素原子におけるシュレーディンガー方程式の解では、ハミルトニアンが

と書ける二粒子系の時間非依存なシュレーディンガー方程式の厳密解を解く。 ～～後略～～
context_section_description :  水素原子の波動関数のl=0～3における角因子は以下のようになる。ここで、はそれぞれ、動径方向の関数 ～～後略～～
```

これは水素原子のシュレーディンガー方程式の動径方向の解の一つを図示した画像
<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/613488/037bfb07-9137-726e-127a-0cbd24e3fa69.png", width=400>
に関するデータで、画像は記事の後半にたくさん貼られているもののうちの一つです。
先ほどの場合との主な違いは以下の通りです：

- `section_title`がnanではなく、画像を含むセクションのタイトルとなっている。
- `hierarchical_section_title`は`page_title`と`section_title`を階層的につなげたものになっている。
- `caption_reference_description`は画像の下に付くキャプションなので、今回の画像では存在しない(nan)。
- `context_section_description`は画像を含むセクションの本文なので、記事の一番上に書いてある本文である`context_page_description`とは異なる。

また、元の記事では`caption_reference_description`と`context_section_description`の部分には数式が書かれていますが、WITデータでは数式は削除されてしまっています。したがって、WITデータセットで学習させても残念ながらシュレーディンガー方程式を解けるモデルは作れません。

このように文章項目の質や画像との関連性の強さはまちまちなので、なかなか使い方が難しそうだという印象を受けました。とは言え、wikipediaから収集したデータなので他のwebページから自動収集したデータセットと比べればある程度の品質は担保されているのだと思います。

## 最後に
このような画像とテキストのペアのデータセットの使い道としては、OpenAIの[CLIP](https://arxiv.org/abs/2103.00020)のようなモデルの日本語版を作ることに興味がありますが、本家のCLIPは4億もの画像・テキストペアを使って学習されており、それと比べればWITデータセットの日本語部分は1/400程度の規模ですし、何よりそれほど大規模な学習は個人レベルでできるものではないでしょう。
CLIPの後続のVision&Languageモデルの[ALBEF](https://arxiv.org/abs/2107.07651)や[BLIP](https://arxiv.org/abs/2201.12086)（いずれもSalesforceの研究）は複数のデータセットを合わせた1400万セットのデータで学習されており、CLIPよりは規模が一桁小さいですが、まだWIT日本語部分は一桁及びません。
より小さなデータセットでCLIPを学習させるという方向の研究では、[CLIP-Lite](https://arxiv.org/abs/2112.07133)が11万8千枚の画像と１画像あたり５つの文章から成るCOCO Captionsデータセットで学習したそうで、これはWIT日本語部分よりも小さなデータセットです。ただし、COCOのキャプションは人間がアノテーションしたものなので、webから収集したWITのようなデータよりも画像とテキストの関係性が直接的なものだという違いがあります。その意味では、[STAIR Captions](http://captions.stair.center/)という日本語版のCOCOデータセット（画像16万枚）が使えるのではないかと思っています。