File size: 5,740 Bytes
3afe23e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import streamlit as st
import itertools
from typing import Dict, Union
from nltk import sent_tokenize
import nltk
nltk.download('punkt')
import torch
from transformers import(
AutoModelForSeq2SeqLM,
AutoTokenizer
)
class QGPipeline:
def __init__(
self
):
self.model = AutoModelForSeq2SeqLM.from_pretrained("muchad/idt5-qa-qg")
self.tokenizer = AutoTokenizer.from_pretrained("muchad/idt5-qa-qg")
self.qg_format = "highlight"
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model.to(self.device)
self.ans_model = self.model
self.ans_tokenizer = self.tokenizer
assert self.model.__class__.__name__ in ["T5ForConditionalGeneration"]
self.model_type = "t5"
def __call__(self, inputs: str):
inputs = " ".join(inputs.split())
sents, answers = self._extract_answers(inputs)
flat_answers = list(itertools.chain(*answers))
if len(flat_answers) == 0:
return []
qg_examples = self._prepare_inputs_for_qg_from_answers_hl(sents, answers)
qg_inputs = [example['source_text'] for example in qg_examples]
questions = self._generate_questions(qg_inputs)
output = [{'answer': example['answer'], 'question': que} for example, que in zip(qg_examples, questions)]
return output
def _generate_questions(self, inputs):
inputs = self._tokenize(inputs, padding=True, truncation=True)
outs = self.model.generate(
input_ids=inputs['input_ids'].to(self.device),
attention_mask=inputs['attention_mask'].to(self.device),
max_length=80,
num_beams=4,
)
questions = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in outs]
return questions
def _extract_answers(self, context):
sents, inputs = self._prepare_inputs_for_ans_extraction(context)
inputs = self._tokenize(inputs, padding=True, truncation=True)
outs = self.ans_model.generate(
input_ids=inputs['input_ids'].to(self.device),
attention_mask=inputs['attention_mask'].to(self.device),
max_length=80,
)
dec = [self.ans_tokenizer.decode(ids, skip_special_tokens=True) for ids in outs]
answers = [item.split('<sep>') for item in dec]
answers = [i[:-1] for i in answers]
return sents, answers
def _tokenize(self,
inputs,
padding=True,
truncation=True,
add_special_tokens=True,
max_length=512
):
inputs = self.tokenizer.batch_encode_plus(
inputs,
max_length=max_length,
add_special_tokens=add_special_tokens,
truncation=truncation,
padding="max_length" if padding else False,
pad_to_max_length=padding,
return_tensors="pt"
)
return inputs
def _prepare_inputs_for_ans_extraction(self, text):
sents = sent_tokenize(text)
inputs = []
for i in range(len(sents)):
source_text = "extract answers:"
for j, sent in enumerate(sents):
if i == j:
sent = "<hl> %s <hl>" % sent
source_text = "%s %s" % (source_text, sent)
source_text = source_text.strip()
source_text = source_text + " </s>"
inputs.append(source_text)
return sents, inputs
def _prepare_inputs_for_qg_from_answers_hl(self, sents, answers):
inputs = []
for i, answer in enumerate(answers):
if len(answer) == 0: continue
for answer_text in answer:
sent = sents[i]
sents_copy = sents[:]
answer_text = answer_text.strip()
try:
ans_start_idx = sent.index(answer_text)
sent = f"{sent[:ans_start_idx]} <hl> {answer_text} <hl> {sent[ans_start_idx + len(answer_text): ]}"
sents_copy[i] = sent
source_text = " ".join(sents_copy)
source_text = f"generate question: {source_text}"
if self.model_type == "t5":
source_text = source_text + " </s>"
except:
continue
inputs.append({"answer": answer_text, "source_text": source_text})
return inputs
class TaskPipeline(QGPipeline):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def __call__(self, inputs: Union[Dict, str]):
return super().__call__(inputs)
def pipeline():
task = TaskPipeline
return task()
@st.cache(ttl=24*3600,allow_output_mutation=True)
def pipeline():
task = TaskPipeline
return task()
st.title("Indonesian Question Generation")
st.write("Indonesian Question Generation System using [idT5](https://huggingface.co/muchad/idt5-base)")
qg = pipeline()
default_context = "Kapitan Pattimura adalah pahlawan dari Maluku. Beliau lahir pada tanggal 8 Juni 1783 dan meninggal pada tanggal 16 Desember 1817."
context_in = st.text_area('Context:', default_context, height=200)
if st.button('Generate Question'):
if context_in:
questions = qg(context_in)
re = ""
for i, q in enumerate(questions):
re += (str(i+1) + "\tAnswer: %s".expandtabs(1) % q['answer'] + " \n" + "\tQuestion: %s".expandtabs(2) % q['question'] + " \n")
st.write(re)
else:
st.write("Please check your context")
|