Electraを用いてチャットボットを作ってみた

こんにちにゃんです。
水色桜（みずいろさくら）です。
今回はBERTの改良版であるElectraを用いて英語のチャットボットを作っていこうと思います。
本記事はBERT後継モデル？BERTにGANの枠組みを取り入れたElectraが、RoBERTaの1/4の学習データで、RoBERTaと同じ精度を達成！Electraの仕組みを徹底解説を参考にさせていただきました。
またモデルとしてelectra-base-squad2を使わせていただきました。感謝いたします。

Electra

２０１９年９月２６日にgoogleさんが公開した学習済みモデルです。RoBERTaの1/4の計算量で、RoBERTaと同等の精度を達成しました。
electraはgeneratorとdiscriminatorを用いて事前学習を行っており、GAN(Generative Adversarial Network; 敵対的生成ネットワーク)のようなアイデアを自然言語処理の事前学習に流用した画期的な研究です。
BERTは文章中の15％の単語を[MASK]に置き換え、この部分の元の単語は何だったかを予測することでモデルの事前学習を行うため、文章の１５％しか学習に用いることができないという課題がありました。
それに対してElectraではgeneratorが文章中の15％の単語を別の単語に置き換え、discriminatorが置き換えられた単語かどうかを学んでいきます。それにより文章のすべての単語を予測することになるため、文章中の全単語について学習が行えます。

詳しくはBERT後継モデル？BERTにGANの枠組みを取り入れたElectraが、RoBERTaの1/4の学習データで、RoBERTaと同じ精度を達成！Electraの仕組みを徹底解説をご覧いただければと思います。本記事では説明を省略させていただきます。

実装

まずmodelとtokenizerのロードを行います。モデルはelectra-base-squad2を使わせていただきました。感謝いたします。初回時に自動的にダウンロードされます。容量が500M弱あるので注意してください。

model_load.py

# モデルの選択（初回時に自動的にダウンロードされます）
model_name = "deepset/electra-base-squad2"

# modelとtokenizerのロード
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)

次に返答を返す関数を定義します。contextととなる文章はelectra_qa.txtというファイルに保存しています。ソースコードを使う際はファイル名を変更してください。

answer.py

# テキストファイルからcontextを取得
with open('electra_qa.txt', encoding = 'utf-8') as f:  # 使用する際はファイル名を変えてください 
    text = f.read()
    context = text.replace('\n', '')


# ボタンが押されるとこの関数を実行する
def electra():
    question = textF.get() # テキストボックスから質問を取得
    msgs.insert('end', 'you : '+question) # 質問を表示
    textF.delete(0, 'end') # テキストボックスの中身を削除
    QA_input = {
    'question': question,
    'context': context
    }
    res = nlp(QA_input)
    # resの形式：{'score' : 0.~, 'start' : ~, 'end' : ~, 'answer' : '~'}
    res_list = str(res).split(',')  # ,で区切って配列にする
    res_alf = re.sub(r"[^a-zA-Z]+", " ", res_list[3])  # 配列の内答えに該当する部分からアルファベットのみを抽出する
    RES = res_alf.replace('answer', '')  # answerという文字を削除する
    msgs.insert('end', 'system :'+RES)  # 

最後にソースコード全体を載せます。もしお忙しい方はコピペして使ってみてください。

electra.py

# 必要なモジュールをインポート
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
import re
import tkinter as tk


# tkinterを用いてGUIを作成
root = tk.Tk()  # ウィンドウを作成
root.title(u'Question-Answering system by electra')  # タイトルの定義
root.geometry('520x480')  # ウィンドウサイズを定義
frame = tk.Frame(root, bg='Green yellow')  # テキストボックスなどを載せるフレームを定義
frame.pack()  # フレームを設置
sc = tk.Scrollbar(frame)  # スクロールバーの定義
sc.pack(side = 'right', fill = 'y')  # スクロールバーを設置
msgs = tk.Listbox(frame,width = 70, height = 24, x = 0, y = 0, yscrollcommand = sc.set, bg = 'azure', fg = 'black')  # テキストボックスの定義
msgs.pack(side = 'left', fill = 'both', pady = 20)  # テキストボックスの設置
msgs.insert('end', 'system : Please give me a question')  # テキストの表示

# モデルの選択（初回時に自動的にダウンロードされます）
model_name = "deepset/electra-base-squad2"

# modelとtokenizerのロード
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)


# テキストファイルからcontextを取得
with open('electra_qa.txt', encoding = 'utf-8') as f:  # 使用する際はファイル名を変えてください 
    text = f.read()
    context = text.replace('\n', '')


# ボタンが押されるとこの関数を実行する
def electra():
    question = textF.get() # テキストボックスから質問を取得
    msgs.insert('end', 'you : '+question) # 質問を表示
    textF.delete(0, 'end') # テキストボックスの中身を削除
    QA_input = {
    'question': question,
    'context': context
    }
    res = nlp(QA_input)
    # resの形式：{'score' : 0.~, 'start' : ~, 'end' : ~, 'answer' : '~'}
    res_list = str(res).split(',')  # ,で区切って配列にする
    res_alf = re.sub(r"[^a-zA-Z]+", " ", res_list[3])  # 配列の内答えに該当する部分からアルファベットのみを抽出する
    RES = res_alf.replace('answer', '')  # answerという文字を削除する
    msgs.insert('end', 'system :'+RES)  # 


#ボタンとテキストボックスの定義
btn = tk.Button(root, text = '送信', font = ('utf-8_sig', 10), bg = 'cyan', command = electra)
btn.place(x = 480,y = 435)
textF = tk.Entry(root, font = ('utf-8_sig', 15), width = 40)
textF.place(x = 70,y = 435)
label2 = tk.Label(root, text = 'question', font = ('utf-8_sig', 10))
label2.place(x = 10,y = 435)

root.mainloop()

実行するとこのようになります。適切に返答を返してくれます。

いかがでしたでしょうか？
比較的簡単にチャットボットが作れたように感じます。
使用した感じはBERTとあまり変わらないように感じました。
BERTのほうが解説記事は圧倒的に多いので、あまりElectraに関する体験談がないのも、頷ける気がします。
今回はここまでにしようと思います。
では、ばいにゃん～。