Add application file

app.py

@@ -0,0 +1,39 @@
+import gradio as gr
+import random
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from question_generation import question_generation_sampling
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+g1_tokenizer = AutoTokenizer.from_pretrained("potsawee/t5-large-generation-squad-QuestionAnswer")
+g1_model = AutoModelForSeq2SeqLM.from_pretrained("potsawee/t5-large-generation-squad-QuestionAnswer")
+g2_tokenizer = AutoTokenizer.from_pretrained("potsawee/t5-large-generation-race-Distractor")
+g2_model = AutoModelForSeq2SeqLM.from_pretrained("potsawee/t5-large-generation-race-Distractor")
+g1_model.eval()
+g2_model.eval()
+g1_model.to(device)
+g2_model.to(device)
+
+
+def generate_multiple_choice_question(
+    context
+):
+    num_questions = 1
+    question_item = question_generation_sampling(
+        g1_model, g1_tokenizer,
+        g2_model, g2_tokenizer,
+        context, num_questions, device
+    )[0]
+    question = question_item['question']
+    options = question_item['options']
+    options[0] = f"{options[0]} [ANSWER]"
+    random.shuffle(options)
+    output_string = f"Question: {question}\n[A] {options[0]}\n[B] {options[1]}\n[C] {options[2]}\n[D] {options[3]}"
+    return output_string
+
+demo = gr.Interface(
+    fn=generate_multiple_choice_question,
+    inputs=gr.Textbox(lines=5, placeholder="Context Here..."),
+    outputs=gr.Textbox(lines=5, placeholder="Question: ...\n[A] ...\n[B] ...\n[C] ...\n[D] ..."),
+)
+demo.launch()

question_generation.py

@@ -0,0 +1,97 @@
+import torch
+import re
+
+@torch.no_grad()
+def question_generation_sampling(
+    g1_model,
+    g1_tokenizer,
+    g2_model,
+    g2_tokenizer,
+    context,
+    num_questions,
+    device,
+):
+    qa_input_ids = prepare_qa_input(
+            g1_tokenizer,
+            context=context,
+            device=device,
+    )
+    max_repeated_sampling = int(num_questions * 1.5) # sometimes generated question+answer is invalid
+    num_valid_questions = 0
+    questions = []
+    for q_ in range(max_repeated_sampling):
+        # Stage G.1: question+answer generation
+        outputs = g1_model.generate(
+            qa_input_ids,
+            max_new_tokens=128,
+            do_sample=True,
+        )
+        question_answer = g1_tokenizer.decode(outputs[0], skip_special_tokens=False)
+        question_answer = question_answer.replace(g1_tokenizer.pad_token, "").replace(g1_tokenizer.eos_token, "")
+        question_answer_split = question_answer.split(g1_tokenizer.sep_token)
+        if len(question_answer_split) == 2:
+            # valid Question + Annswer output
+            num_valid_questions += 1
+        else:
+            continue
+        question = question_answer_split[0].strip()
+        answer = question_answer_split[1].strip()
+
+        # Stage G.2: Distractor Generation
+        distractor_input_ids = prepare_distractor_input(
+            g2_tokenizer,
+            context = context,
+            question = question,
+            answer = answer,
+            device = device,
+            separator = g2_tokenizer.sep_token,
+        )
+        outputs = g2_model.generate(
+            distractor_input_ids,
+            max_new_tokens=128,
+            do_sample=True,
+        )
+        distractors = g2_tokenizer.decode(outputs[0], skip_special_tokens=False)
+        distractors = distractors.replace(g2_tokenizer.pad_token, "").replace(g2_tokenizer.eos_token, "")
+        distractors = re.sub("<extra\S+>", g2_tokenizer.sep_token, distractors)
+        distractors = [y.strip() for y in distractors.split(g2_tokenizer.sep_token)]
+        options = [answer] + distractors
+
+        while len(options) < 4:
+            options.append(options[-1])
+
+        question_item = {
+            'question': question,
+            'options': options,
+        }
+        questions.append(question_item)
+        if num_valid_questions == num_questions:
+            break
+    return questions
+
+
+def prepare_qa_input(t5_tokenizer, context, device):
+    """
+    input: context
+    output: question <sep> answer
+    """
+    encoding = t5_tokenizer(
+        [context],
+        return_tensors="pt",
+    )
+    input_ids = encoding.input_ids.to(device)
+    return input_ids
+
+
+def prepare_distractor_input(t5_tokenizer, context, question, answer, device, separator='<sep>'):
+    """
+    input: question <sep> answer <sep> article
+    output: distractor1 <sep> distractor2 <sep> distractor3
+    """
+    input_text = question + ' ' + separator + ' ' + answer + ' ' + separator + ' ' + context
+    encoding = t5_tokenizer(
+        [input_text],
+        return_tensors="pt",
+    )
+    input_ids = encoding.input_ids.to(device)
+    return input_ids

requirements.txt

@@ -0,0 +1,2 @@
+torch>=1.10
+transformers>=4.11.3