LLM(大規模言語モデル)に医師国家試験を解かせる

ここではCyberAgentが開発した商用利用可能なLLMであるCALMを使って医師国家試験を解かせる方法を紹介します。

Google Colabの無料版でも動作する1Bのモデルを使ったものになっているので、どんな感じで使えばいいのかを知りたい人におすすめです。

また、もっと性能の良い7Bのモデルや他のLLMもモデル定義まわりだけを書き換えれば使えるので、本例を練習として使い、本番のクラウドサービス上では上位モデルを使ってより良い結果を出したりモデル間比較をするのも良いかもしれません。

なお以下の例では医師国家試験のデータとして、"IgakuQA"のものを利用しています。ただ、IgakuQAはテキストデータのみであり、図表は含まれていないということには注意してください。テキストだけでなく画像もLLMのinputとして渡したい場合は自分で拡張するか、他のデータセットを使う必要があります。

お持ちのデータを利用する場合はファイルの読み込みの箇所を変えるといったように適宜修正してお使いください。

必要なパッケージのインストール

!pip install jsonlines

!pip install transformers

!pip install tqdm

医師国家試験データ(IgakuQA)のダウンロード

!git clone https://github.com/jungokasai/IgakuQA.git

使用する関数の定義
(一部のものはIgakuQAのものを改変)

import string

# LLMへのinput(prompt)をつくる
def create_input(prompt, question):
    messages = ["医師国家試験を解きます。\n"]
    for example in prompt:
        messages.extend(dict2problem(example))
    messages.extend(dict2problem(question, False))
    return messages

# 問題文の整形
def dict2problem(dict_input, demo=True):
    problem = "問題: " + dict_input['problem_text']
    choices = dict_input['choices']
    answer = dict_input['answer']
    if len(choices) > 0:
        for choice, label in zip(choices, string.ascii_lowercase):
            problem = problem + '\n' + label + ': ' + choice
        problem = problem + "\n必ずa,b,c,d,eの中からちょうど{}個選んでください。".format(len(answer))
        problem = problem + "\n答え:"
    output = [problem]
    if not demo:
        return output
    output.append(",".join(answer)+"\n")
    return output

# LLMから結果を受け取る
def ask_llm(input):
    inputs = tokenizer(input, return_tensors="pt").to(model.device)
    with torch.no_grad():
        tokens = model.generate(
            **inputs,
            max_new_tokens=64,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id,
        )
    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
    return output

LLMの定義

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 無料版のColabだと3bや7bは無理
model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-1b")
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-1b")

LLMに医師国家試験を解かせる

以下の例では2022年のA問題を使用。
Colabの無料版CPUモードでだいたい22分かかります。

from IgakuQA.scripts.utils.tools import read_jsonl
import tqdm

# 問題ファイルの指定
in_file = "IgakuQA/data/2022/116-A.jsonl"

# promptファイルの指定(以下の例では使っていない)
# promptファイルの中身は問題文と解答の例が3問入っている
prompt_file = "IgakuQA/scripts/prompts/prompt.jsonl"

# ファイル読み込み
questions = read_jsonl(in_file)
prompt = []# promptファイルのデータを使う場合はread_jsonl(prompt_file)を指定する

ids = []
preds = []
answers = []
for question in tqdm.tqdm(questions):
    #print(question['problem_id'])
    llm_input = create_input(prompt, question)
    res = ask_llm("".join(llm_input))
    res = res.split("答え:")[-1]
    ids.append(question['problem_id'])
    preds.append(res)
    answers.append(",".join(question["answer"]))

LLMの答えを整形

LLMが返してきたものには不要な文字が含まれている場合があるので、それらを除去します。

以下の例でも除去できない場合もあると思いますが、こういった処理はよく使うものなので、この機会に練習もかねて習得しておきましょう。

def reshape_ans(input):
    reshaped = input.strip()
    reshaped = reshaped.split("(")[0]
    reshaped = reshaped.split("です")[0]
    reshaped = reshaped.split("は")[0]
    reshaped = reshaped.split("に")[0]
    reshaped = reshaped.split(" ")[0]
    reshaped = reshaped.replace("、", ",")
    reshaped = reshaped.lower()
    return reshaped

preds_reshaped = []
for k in preds:
    k = reshape_ans(k)
    preds_reshaped.append(k)
    #print(k)

正解率の計算

ここでは簡単に正解率だけを計算します。

n = 0
for a, b in zip(answers, preds_reshaped):
    if a == b:
        n += 1

print("Accuracy: ", n/len(answers))

結果を保存することで、IgakuQAのscriptを使って評価することもできます。

from IgakuQA.scripts.utils.tools import answer2jsonl

out_file = "116-A_open-calm-1b.jsonl"

answer2jsonl(preds_reshaped, preds_reshaped, questions, out_file)
# answersとpreds_reshapedの間違いと思われるかもしれませんが、以下のevaluate_main.pyの挙動をみるとこれが正しいようです。

!python IgakuQA/scripts/evaluate_main.py --gold-file IgakuQA/data/2022/116-A.jsonl --pred-file 116-A_open-calm-1b.jsonl

比較対象として、IgakuQAに内包されている結果の中からChatGPTを使った場合のものを表示させます。

!python IgakuQA/scripts/evaluate_main.py --gold-file IgakuQA/data/2022/116-A.jsonl --pred-file IgakuQA/baseline_results/2022/116-A_chatgpt.jsonl

やはりChatGPTと比べると見劣りしてしまいますが、ChatGPTのパラメータ数が少なくとも数千億(今回使ったCALM-1bの100倍以上)と言われていることを考えると、結構がんばっているほうなのではないでしょうか。

また、無料版Colabでは厳しいですが、より上位の7bモデルを使ったり、過去問データでFine-Tuningしてあげることでもっと良い結果が出るようになる可能性があると思います。

以上のコードがすでに記入されたnotebook(.ipynb)を置いておきます。
Google Colabにアップロードしてすぐ使えるので、よろしければご利用ください。

配布notebookの使い方についてはこちら

他にもこういった記事が読みたい！これを解説してほしい！などご意見がありましたら、アンケートフォームからご意見をいただけるとうれしいです。