sunwaee/MT5-Questions-Answers-Generation-Extraction
16
1import os2 3import gdown as gdown4import nltk5import streamlit as st6import torch7from transformers import AutoTokenizer8 9from mt5 import MT510 11 12def download_models(ids):13 """14 Download all models.15 16 :param ids: name and links of models17 :return:18 """19 20 # Download sentence tokenizer21 nltk.download('punkt')22 23 # Download model from drive if not stored locally24 for key in ids:25 if not os.path.isfile(f"model/{key}.ckpt"):26 url = f"https://drive.google.com/u/0/uc?id={ids[key]}"27 gdown.download(url=url, output=f"model/{key}.ckpt")28 29 30@st.cache(allow_output_mutation=True)31def load_model(model_path):32 """33 Load model and cache it.34 35 :param model_path: path to model36 :return:37 """38 39 device = 'cuda' if torch.cuda.is_available() else 'cpu'40 41 # Loading model and tokenizer42 model = MT5.load_from_checkpoint(model_path).eval().to(device)43 model.tokenizer = AutoTokenizer.from_pretrained('tokenizer')44 45 return model46 47 48# Page config49st.set_page_config(layout="centered")50st.title("Questions/Answers Pairs Gen.")51st.write("Question Generation, Question Answering and Questions/Answers Generation using Google MT5. ")52 53# Variables54ids = {'mt5-small': st.secrets['small'],55 'mt5-base': st.secrets['base']}56 57 58# Download all models from drive59download_models(ids)60 61# Task selection62 63left, right = st.columns([4, 2])64task = left.selectbox('Choose the task: ',65 options=['Questions/Answers Pairs Generation', 'Question Answering', 'Question Generation'],66 help='Choose the task you want to try out')67 68# Model selection69model_path = right.selectbox('', options=[k for k in ids], index=1, help='Model to use. ')70model = load_model(model_path=f"model/{model_path}.ckpt")71right.write(model.device)72 73if task == 'Questions/Answers Pairs Generation':74 # Input area75 inputs = st.text_area('Context:', value="A few years after the First Crusade, in 1107, the Normans under "76 "the command of Bohemond, Robert\'s son, landed in Valona and "77 "besieged Dyrrachium using the most sophisticated military "78 "equipment of the time, but to no avail. Meanwhile, they occupied "79 "Petrela, the citadel of Mili at the banks of the river Deabolis, "80 "Gllavenica (Ballsh), Kanina and Jericho. This time, "81 "the Albanians sided with the Normans, dissatisfied by the heavy "82 "taxes the Byzantines had imposed upon them. With their help, "83 "the Normans secured the Arbanon passes and opened their way to "84 "Dibra. The lack of supplies, disease and Byzantine resistance "85 "forced Bohemond to retreat from his campaign and sign a peace "86 "treaty with the Byzantines in the city of Deabolis. ", max_chars=2048,87 height=250)88 split = st.checkbox('Split into sentences', value=True)89 90 if split:91 # Split into sentences92 sent_tokenized = nltk.sent_tokenize(inputs)93 res = {}94 95 with st.spinner('Please wait while the inputs are being processed...'):96 # Iterate over sentences97 for sentence in sent_tokenized:98 predictions = model.multitask([sentence], max_length=512)99 questions, answers, answers_bis = predictions['questions'], predictions['answers'], predictions[100 'answers_bis']101 102 # Build answer dict103 content = {}104 for question, answer, answer_bis in zip(questions[0], answers[0], answers_bis[0]):105 content[question] = {'answer (extracted)': answer, 'answer (generated)': answer_bis}106 res[sentence] = content107 108 # Answer area109 st.write(res)110 111 else:112 with st.spinner('Please wait while the inputs are being processed...'):113 # Prediction114 predictions = model.multitask([inputs], max_length=512)115 questions, answers, answers_bis = predictions['questions'], predictions['answers'], predictions[116 'answers_bis']117 118 # Answer area119 zip = zip(questions[0], answers[0], answers_bis[0])120 content = {}121 for question, answer, answer_bis in zip:122 content[question] = {'answer (extracted)': answer, 'answer (generated)': answer_bis}123 124 st.write(content)125 126elif task == 'Question Answering':127 128 # Input area129 inputs = st.text_area('Context:', value="A few years after the First Crusade, in 1107, the Normans under "130 "the command of Bohemond, Robert\'s son, landed in Valona and "131 "besieged Dyrrachium using the most sophisticated military "132 "equipment of the time, but to no avail. Meanwhile, they occupied "133 "Petrela, the citadel of Mili at the banks of the river Deabolis, "134 "Gllavenica (Ballsh), Kanina and Jericho. This time, "135 "the Albanians sided with the Normans, dissatisfied by the heavy "136 "taxes the Byzantines had imposed upon them. With their help, "137 "the Normans secured the Arbanon passes and opened their way to "138 "Dibra. The lack of supplies, disease and Byzantine resistance "139 "forced Bohemond to retreat from his campaign and sign a peace "140 "treaty with the Byzantines in the city of Deabolis. ", max_chars=2048,141 height=250)142 question = st.text_input('Question:', value="What forced Bohemond to retreat from his campaign? ")143 144 # Prediction145 with st.spinner('Please wait while the inputs are being processed...'):146 predictions = model.qa([{'question': question, 'context': inputs}], max_length=512)147 answer = {question: predictions[0]}148 149 # Answer area150 st.write(answer)151 152elif task == 'Question Generation':153 154 # Input area155 inputs = st.text_area('Context (highlight answers with <hl> tokens): ',156 value="A few years after the First Crusade, in <hl> 1107 <hl>, the <hl> Normans <hl> under "157 "the command of <hl> Bohemond <hl>, Robert\'s son, landed in Valona and "158 "besieged Dyrrachium using the most sophisticated military "159 "equipment of the time, but to no avail. Meanwhile, they occupied "160 "Petrela, <hl> the citadel of Mili <hl> at the banks of the river Deabolis, "161 "Gllavenica (Ballsh), Kanina and Jericho. This time, "162 "the Albanians sided with the Normans, dissatisfied by the heavy "163 "taxes the Byzantines had imposed upon them. With their help, "164 "the Normans secured the Arbanon passes and opened their way to "165 "Dibra. The <hl> lack of supplies, disease and Byzantine resistance <hl> "166 "forced Bohemond to retreat from his campaign and sign a peace "167 "treaty with the Byzantines in the city of Deabolis. ", max_chars=2048,168 height=250)169 170 # Split by highlights171 hl_index = [i for i in range(len(inputs)) if inputs.startswith('<hl>', i)]172 contexts = []173 answers = []174 175 # Build a context for each highlight pair176 for i in range(0, len(hl_index), 2):177 contexts.append(inputs[:hl_index[i]].replace('<hl>', '') +178 inputs[hl_index[i]: hl_index[i + 1] + 4] +179 inputs[hl_index[i + 1] + 4:].replace('<hl>', ''))180 answers.append(inputs[hl_index[i]: hl_index[i + 1] + 4].replace('<hl>', '').strip())181 182 # Prediction183 with st.spinner('Please wait while the inputs are being processed...'):184 predictions = model.qg(contexts, max_length=512)185 186 # Answer area187 content = {}188 for pred, ans in zip(predictions, answers):189 content[pred] = ans190 st.write(content)191