Chapter5.ipynb

{
    "nbformat": 4,
    "nbformat_minor": 0,
    "metadata": {
        "accelerator": "GPU",
        "colab": {
            "name": "Chapter05.ipynb",
            "provenance": [],
            "collapsed_sections": []
        },
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.7.9"
        }
    },
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "QB3WEyIVstl0"
            },
            "source": [
                "# 5章\n",
                "- 以下で実行するコードには確率的な処理が含まれていることがあり、コードの出力結果と本書に記載されている出力例が異なることがあります。"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "kvqSUAEtU_VJ"
            },
            "source": [
                "# 5-1\n",
                "!pip install transformers==4.18.0 fugashi==1.1.0 ipadic==1.0.0"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "DWT32lOgHLrU"
            },
            "source": [
                "# 5-2\n",
                "import numpy as np\n",
                "import torch\n",
                "from transformers import BertJapaneseTokenizer, BertForMaskedLM"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "I7X-Iy52AC1v"
            },
            "source": [
                "# 5-3\n",
                "model_name = 'tohoku-nlp/bert-base-japanese-whole-word-masking'\n",
                "tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)\n",
                "bert_mlm = BertForMaskedLM.from_pretrained(model_name)\n",
                "bert_mlm = bert_mlm.cuda()"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "EfKt-j0WLOfx"
            },
            "source": [
                "# 5-4\n",
                "text = '今日は[MASK]へ行く。'\n",
                "tokens = tokenizer.tokenize(text)\n",
                "print(tokens)"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "YaW5Y9fM5zeM"
            },
            "source": [
                "# 5-5\n",
                "# 文章を符号化し、GPUに配置する。\n",
                "input_ids = tokenizer.encode(text, return_tensors='pt')\n",
                "input_ids = input_ids.cuda()\n",
                "\n",
                "# BERTに入力し、分類スコアを得る。\n",
                "# 系列長を揃える必要がないので、単にiput_idsのみを入力します。\n",
                "with torch.no_grad():\n",
                "    output = bert_mlm(input_ids=input_ids)\n",
                "    scores = output.logits"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "Z-5lnX9r8XKl"
            },
            "source": [
                "# 5-6\n",
                "# ID列で'[MASK]'(IDは4)の位置を調べる\n",
                "mask_position = input_ids[0].tolist().index(4) \n",
                "\n",
                "# スコアが最も良いトークンのIDを取り出し、トークンに変換する。\n",
                "id_best = scores[0, mask_position].argmax(-1).item()\n",
                "token_best = tokenizer.convert_ids_to_tokens(id_best)\n",
                "token_best = token_best.replace('##', '')\n",
                "\n",
                "# [MASK]を上で求めたトークンで置き換える。\n",
                "text = text.replace('[MASK]',token_best)\n",
                "\n",
                "print(text)"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "TgbIA-1-EVaJ"
            },
            "source": [
                "# 5-7\n",
                "def predict_mask_topk(text, tokenizer, bert_mlm, num_topk):\n",
                "    \"\"\"\n",
                "    文章中の最初の[MASK]をスコアの上位のトークンに置き換える。\n",
                "    上位何位まで使うかは、num_topkで指定。\n",
                "    出力は穴埋めされた文章のリストと、置き換えられたトークンのスコアのリスト。\n",
                "    \"\"\"\n",
                "    # 文章を符号化し、BERTで分類スコアを得る。\n",
                "    input_ids = tokenizer.encode(text, return_tensors='pt')\n",
                "    input_ids = input_ids.cuda()\n",
                "    with torch.no_grad():\n",
                "        output = bert_mlm(input_ids=input_ids)\n",
                "    scores = output.logits\n",
                "\n",
                "    # スコアが上位のトークンとスコアを求める。\n",
                "    mask_position = input_ids[0].tolist().index(4) \n",
                "    topk = scores[0, mask_position].topk(num_topk)\n",
                "    ids_topk = topk.indices # トークンのID\n",
                "    tokens_topk = tokenizer.convert_ids_to_tokens(ids_topk) # トークン\n",
                "    scores_topk = topk.values.cpu().numpy() # スコア\n",
                "\n",
                "    # 文章中の[MASK]を上で求めたトークンで置き換える。\n",
                "    text_topk = [] # 穴埋めされたテキストを追加する。\n",
                "    for token in tokens_topk:\n",
                "        token = token.replace('##', '')\n",
                "        text_topk.append(text.replace('[MASK]', token, 1))\n",
                "\n",
                "    return text_topk, scores_topk\n",
                "\n",
                "text = '今日は[MASK]へ行く。'\n",
                "text_topk, _ = predict_mask_topk(text, tokenizer, bert_mlm, 10)\n",
                "print(*text_topk, sep='\\n')"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "yCaGV_rT3A5N"
            },
            "source": [
                "# 5-8\n",
                "def greedy_prediction(text, tokenizer, bert_mlm):\n",
                "    \"\"\"\n",
                "    [MASK]を含む文章を入力として、貪欲法で穴埋めを行った文章を出力する。\n",
                "    \"\"\"\n",
                "    # 前から順に[MASK]を一つづつ、スコアの最も高いトークンに置き換える。\n",
                "    for _ in range(text.count('[MASK]')):\n",
                "        text = predict_mask_topk(text, tokenizer, bert_mlm, 1)[0][0]\n",
                "    return text\n",
                "\n",
                "text = '今日は[MASK][MASK]へ行く。'\n",
                "greedy_prediction(text, tokenizer, bert_mlm)"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "prdEvsxBrrGq"
            },
            "source": [
                "# 5-9\n",
                "text = '今日は[MASK][MASK][MASK][MASK][MASK]'\n",
                "greedy_prediction(text, tokenizer, bert_mlm)"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "yHRemOdN0QE9"
            },
            "source": [
                "# 5-10\n",
                "def beam_search(text, tokenizer, bert_mlm, num_topk):\n",
                "    \"\"\"\n",
                "    ビームサーチで文章の穴埋めを行う。\n",
                "    \"\"\"\n",
                "    num_mask = text.count('[MASK]')\n",
                "    text_topk = [text]\n",
                "    scores_topk = np.array([0])\n",
                "    for _ in range(num_mask):\n",
                "        # 現在得られている、それぞれの文章に対して、\n",
                "        # 最初の[MASK]をスコアが上位のトークンで穴埋めする。\n",
                "        text_candidates = [] # それぞれの文章を穴埋めした結果を追加する。\n",
                "        score_candidates = [] # 穴埋めに使ったトークンのスコアを追加する。\n",
                "        for text_mask, score in zip(text_topk, scores_topk):\n",
                "            text_topk_inner, scores_topk_inner = predict_mask_topk(\n",
                "                text_mask, tokenizer, bert_mlm, num_topk\n",
                "            )\n",
                "            text_candidates.extend(text_topk_inner)\n",
                "            score_candidates.append( score + scores_topk_inner )\n",
                "\n",
                "        # 穴埋めにより生成された文章の中から合計スコアの高いものを選ぶ。\n",
                "        score_candidates = np.hstack(score_candidates)\n",
                "        idx_list = score_candidates.argsort()[::-1][:num_topk]\n",
                "        text_topk = [ text_candidates[idx] for idx in idx_list ]\n",
                "        scores_topk = score_candidates[idx_list]\n",
                "\n",
                "    return text_topk\n",
                "\n",
                "text = \"今日は[MASK][MASK]へ行く。\"\n",
                "text_topk = beam_search(text, tokenizer, bert_mlm, 10)\n",
                "print(*text_topk, sep='\\n')"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "5mhL-VSTvUo7"
            },
            "source": [
                "# 5-11\n",
                "text = '今日は[MASK][MASK][MASK][MASK][MASK]'\n",
                "text_topk = beam_search(text, tokenizer, bert_mlm, 10)\n",
                "print(*text_topk, sep='\\n')"
            ],
            "execution_count": null,
            "outputs": []
        }
    ]
}