app.py

# -*- coding: utf-8 -*-
import gradio as gr
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer, AutoTokenizer, AutoModelForSeq2SeqLM, RobertaForQuestionAnswering

# 0.モデルのロード, Examplesの準備
# TODO 評価対象の要約モデルをロード
tokenizer_sum = AutoTokenizer.from_pretrained("tsmatz/mt5_summarize_japanese")
model_sum = AutoModelForSeq2SeqLM.from_pretrained("tsmatz/mt5_summarize_japanese")

# TODO 質問文の生成モデルをロード
tokenizer_gen_q = T5Tokenizer.from_pretrained("sonoisa/t5-base-japanese-question-generation")
model_gen_q = T5ForConditionalGeneration.from_pretrained("sonoisa/t5-base-japanese-question-generation")

# TODO 回答の生成モデルをロード
tokenizer_qa = AutoTokenizer.from_pretrained("tsmatz/roberta_qa_japanese")
model_qa = RobertaForQuestionAnswering.from_pretrained("tsmatz/roberta_qa_japanese")

# Example 1
eg_text_1 = """
ポケットモンスターの原点は、1996年2月27日に発売されたゲームボーイ用ソフト『ポケットモンスター 赤・緑』である。
開発元はゲームフリーク。コンセプトメーカーにしてディレクターを務めたのは、同社代表取締役でもある田尻智。
この作品が小学生を中心に、口コミから火が点き大ヒットとなり、以降も多くの続編が発売されている（詳しくは「ポケットモンスター（ゲーム）」を参照）。
ゲーム本編作品だけでなく、派生作品や関連作品が数多く発売されている（詳しくはポケットモンスターの関連ゲームを参照）。

ポケモンはゲームのみならず、アニメ化、キャラクター商品化、カードゲーム、アーケードゲームと様々なメディアミックス展開がなされ、日本国外でも人気を獲得している。

ポケモン関連ゲームソフトの累計出荷数は、全世界で2017年11月時点で3億本以上[1]、2022年3月時点で4億4000万本以上に達している[2]。
その中で、メインシリーズの累計販売本数は2016年2月時点での最新作、ニンテンドー3DS『オメガルビー・アルファサファイア』までの25作品で2億100万本となる[3]。
"""
eg_ans_1_1 = "2月27日"
eg_ans_1_2 = "ポケットモンスター 赤・緑"

# Example 2
eg_text_2 = """
アンパンマンの生みの親であるやなせたかしの作品で1968年に「バラの花とジョー」、
「チリンの鈴」の絵本や映画にいち早くアンパンマンが登場しているが、この時はまだ人間の姿。
この童話は一年間連載された。[5]アンパンマン、やなせたかしの作品としての、「アンパンマン」は、
PHP研究所が発行する青年向け雑誌『PHP』の通巻第257号に当たる、『こどものえほん』の1969年10月号[6]（同年10月1日刊行）に掲載された青年向け読物、
やなせたかし（絵と文）「アンパンマン」という形が初出である[7][8][9]。
この時期、やなせが『こどものえほん』のために執筆した読物は連載12本の短編で、「アンパンマン」はその6本目の作品であった。
これら12篇は、株式会社山梨シルクセンター（※3年後、株式会社サンリオへ社名変更）より単行本『十二の真珠』名義で1970年に刊行された。

空腹に喘ぐ人の所へ駆け付けて、自らの大事な持ち物であるパンを差し出して食べるよう勧めるという、のちのアンパンマンに通じる物語の骨組みが、
この作品のおいて早くも整えられている[10][6]。
絵本・漫画・アニメなど、のちに描かれるアンパンマンとの大きな違いと言えば、第一に主人公のアンパンマンが普通の人間のおじさんであり[10][6]、
パンは所有物に過ぎなかったことである。
"""
eg_ans_2_1 = "アンパンマン"
eg_ans_2_2 = "やなせたかし"

"""#### イベント用の関数の実装
- 要約生成の関数作成
- 質問生成の関数作成
- 回答生成の関数作成
"""

# 1. イベント用の関数
def summy(text):
    """要約

    Args
        text: str
            要約対象のテキスト

    Returns
        summarize_text: str
            要約結果のテキスト

    TODO
        処理の実装
    """
    inputs = tokenizer_sum("summarize: " + text, return_tensors="pt")
    outputs = model_sum.generate(
        inputs["input_ids"],
        max_new_tokens=300,
        min_length=150,
        num_beams=5
    )
    summarize_text = tokenizer_sum.decode(outputs[0], skip_special_tokens=True)
    return summarize_text

def generate_questions(answer_1, answer_2, text):
    """質問生成

    Args
        answers: list[str]
            質問生成のための正解単語のリスト
        text: str
            質問文を生成する際に参照するテキスト
        
    Returns
        generated_questions: list[str]
            生成された質問文のリスト

    TODO
        処理の実装
    """
# 質問文の生成
    answer_context_list = [(answer_1, text), (answer_2, text)]

    generated_questions = []

    for answer, context in answer_context_list:
        input = tokenizer_gen_q(f"answer: {answer} context: {context}", return_tensors="pt")

        # 質問文を生成する
        output = model_gen_q.generate(
            input['input_ids'],
            max_new_tokens=100,
            num_beams=4
        )

        # 生成された問題文のトークン列を文字列に変換する。
        output = tokenizer_gen_q.decode(output[0], skip_special_tokens=True)

        generated_questions.append(output)

    return generated_questions

def extract_answer(question, text):
    """質問応答

    Args
        question: str
            質問文のテキスト
        text: str
            質問に回答するために参照するテキスト
    
    Returns
        answer: str
            回答のテキスト

    TODO
        処理の実装
    """
    inputs = tokenizer_qa(question, text, return_tensors="pt")

    outputs = model_qa(**inputs)
    answer_start_scores = outputs.start_logits
    answer_end_scores = outputs.end_logits

    answer_start = torch.argmax(answer_start_scores)
    answer_end = torch.argmax(answer_end_scores) + 1

    input_ids = inputs["input_ids"].tolist()[0]

    answer = tokenizer_qa.decode(input_ids[answer_start:answer_end])

    return answer

def extract_answer_all(gen_q_1, gen_q_2, source_text, sum_text):
    """extract_answer()をまとめて実行する
    TODO
        処理の実装
    """
    a_source_1 = extract_answer(gen_q_1, source_text)
    a_sum_1 = extract_answer(gen_q_1, sum_text)
    a_source_2 = extract_answer(gen_q_2, source_text)
    a_sum_2 = extract_answer(gen_q_2, sum_text)

    return a_source_1, a_sum_1, a_source_2, a_sum_2

"""#### UIの実装
- `gr.Blocks()`を使ったUIの実装
    - 要約生成のUI作成
    - 質問生成のUI作成
    - 回答生成のUI作成
- イベントの実装（btn.clickなど）
    - 要約生成の実行ボタン作成
    - 質問生成の実行ボタン作成
    - 回答生成の実行ボタン作成
"""

# 2. UIの定義
with gr.Blocks() as demo:
    gr.Markdown("### 1. 要約生成")
    # TODO 要約のための入出力UIの作成
    text_source = gr.Textbox(label="要約対象")
    btn_summy = gr.Button("要約生成")
    text_summy = gr.Textbox(label="要約結果")

    gr.Markdown("### 2. 質問生成")
    # TODO 質問文生成のための入力UIの作成
    with gr.Row():
      text_q_1 = gr.Textbox(label="正解1")
      text_q_2 = gr.Textbox(label="正解2")
    btn_generate_questions = gr.Button("質問生成")

    gr.Markdown("### 3. 回答生成")
    # TODO 質問文を表示するUIの作成
    with gr.Row():
      text_gq_1 = gr.Textbox(label="1番目の質問")
      text_gq_2 = gr.Textbox(label="2番目の質問")
    btn_extract_answer = gr.Button("回答生成")
    # TODO それぞれの回答を表示するUIの作成
    with gr.Row():
      with gr.Column():
        text_asrc_1 = gr.Textbox(label="sourceからの答え1")
        text_asum_1 = gr.Textbox(label="sumからの答え1")
      with gr.Column():
        text_asrc_2 = gr.Textbox(label="sourceからの答え2")
        text_asum_2 = gr.Textbox(label="sumからの答え2")

    # 2. イベント発火
    btn_summy.click(
        summy,
        inputs=text_source,  # TODO 定義したUIのコンポーネントを与える
        outputs=text_summy  # TODO 定義したUIのコンポーネントを与える
    )
    btn_generate_questions.click(
        generate_questions,
        inputs=[text_q_1, text_q_2, text_summy],  # TODO 定義したUIのコンポーネントを与える
        outputs=[text_gq_1, text_gq_2]  # TODO 定義したUIのコンポーネントを与える
    )
    btn_extract_answer.click(extract_answer_all,
        inputs=[text_gq_1, text_gq_2, text_source, text_summy],  # TODO 定義したUIのコンポーネントを与える
        outputs=[text_asrc_1, text_asum_1, text_asrc_2, text_asum_2]  # TODO 定義したUIのコンポーネントを与える
    )

    # Examplesの定義
    gr.Markdown("## Examples")
    gr.Examples(
        examples = [[eg_text_1, eg_ans_1_1, eg_ans_1_2], [eg_text_2, eg_ans_2_1, eg_ans_2_2]],
        inputs = [text_source, text_q_1, text_q_2]
    )  # TODO Exampleにデータを与えて、表示させる

"""#### demoの起動"""

demo.launch()