Update pipeline.py
Browse files- pipeline.py +36 -24
pipeline.py
CHANGED
@@ -2,8 +2,8 @@ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
|
|
2 |
from typing import Dict, List, Any
|
3 |
import itertools
|
4 |
from nltk import sent_tokenize
|
|
|
5 |
import nltk
|
6 |
-
import torch
|
7 |
|
8 |
class PreTrainedPipeline():
|
9 |
|
@@ -11,16 +11,20 @@ class PreTrainedPipeline():
|
|
11 |
# IMPLEMENT_THIS
|
12 |
# Preload all the elements you are going to need at inference.
|
13 |
# For instance your model, processors, tokenizer that might be needed.
|
14 |
-
# This function is only called once, so do all the heavy processing I/O here"""
|
15 |
nltk.download('punkt')
|
16 |
-
self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
|
17 |
self.tokenizer = AutoTokenizer.from_pretrained(path)
|
18 |
|
19 |
self.model_type="t5"
|
20 |
-
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
|
|
|
21 |
|
22 |
|
23 |
-
def __call__(self, inputs: str):
|
24 |
if len(inputs) == 0: return []
|
25 |
inputs = " ".join(inputs.split())
|
26 |
sents, answers = self._extract_answers(inputs)
|
@@ -29,16 +33,29 @@ class PreTrainedPipeline():
|
|
29 |
if len(flat_answers) == 0:
|
30 |
return []
|
31 |
|
|
|
|
|
|
|
|
|
|
|
|
|
32 |
qg_examples = self._prepare_inputs_for_qg_from_answers_hl(sents, answers)
|
33 |
-
|
34 |
qg_inputs = [example['source_text'] for example in qg_examples]
|
35 |
questions = self._generate_questions(qg_inputs)
|
36 |
-
|
37 |
-
|
38 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
39 |
|
40 |
def _extract_answers(self, context):
|
41 |
-
print("_extract_answers")
|
42 |
sents, inputs = self._prepare_inputs_for_ans_extraction(context)
|
43 |
inputs = self._tokenize(inputs, padding=True, truncation=True)
|
44 |
|
@@ -50,13 +67,14 @@ class PreTrainedPipeline():
|
|
50 |
|
51 |
dec = [self.tokenizer.decode(ids, skip_special_tokens=False) for ids in outs]
|
52 |
answers = [item.split('<sep>') for item in dec]
|
53 |
-
|
54 |
-
|
|
|
55 |
return sents, answers
|
56 |
|
|
|
57 |
|
58 |
def _prepare_inputs_for_ans_extraction(self, text):
|
59 |
-
print("_prepare_inputs_for_ans_extraction")
|
60 |
sents = sent_tokenize(text)
|
61 |
|
62 |
inputs = []
|
@@ -93,7 +111,6 @@ class PreTrainedPipeline():
|
|
93 |
return inputs
|
94 |
|
95 |
def _generate_questions(self, inputs):
|
96 |
-
print("_generate_questions")
|
97 |
inputs = self._tokenize(inputs, padding=True, truncation=True)
|
98 |
|
99 |
outs = self.model.generate(
|
@@ -105,11 +122,8 @@ class PreTrainedPipeline():
|
|
105 |
|
106 |
questions = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in outs]
|
107 |
return questions
|
108 |
-
|
109 |
-
|
110 |
|
111 |
def _prepare_inputs_for_qg_from_answers_hl(self, sents, answers):
|
112 |
-
print("_prepare_inputs_for_qg_from_answers_hl")
|
113 |
inputs = []
|
114 |
for i, answer in enumerate(answers):
|
115 |
if len(answer) == 0: continue
|
@@ -118,8 +132,6 @@ class PreTrainedPipeline():
|
|
118 |
sents_copy = sents[:]
|
119 |
answer_text = self.remove_pad(answer_text)
|
120 |
answer_text = answer_text.strip()
|
121 |
-
print("Answer", answer)
|
122 |
-
print("Answer text", answer_text)
|
123 |
|
124 |
try:
|
125 |
ans_start_idx = sent.lower().index(answer_text.lower())
|
@@ -139,13 +151,14 @@ class PreTrainedPipeline():
|
|
139 |
|
140 |
return inputs
|
141 |
|
142 |
-
def clean_generated_QAs(self, generated_QAs):
|
143 |
clean_QAs = []
|
144 |
answers_used = set()
|
145 |
# Only allow 1 question per answer, take the first case of it
|
146 |
for qa in generated_QAs:
|
147 |
-
|
148 |
-
|
|
|
149 |
answers_used.add(qa['answer'])
|
150 |
clean_QAs.append(qa)
|
151 |
return clean_QAs
|
@@ -153,5 +166,4 @@ class PreTrainedPipeline():
|
|
153 |
def remove_pad(self, str):
|
154 |
if "<pad>" in str:
|
155 |
return str.replace("<pad>", "")
|
156 |
-
return str
|
157 |
-
|
2 |
from typing import Dict, List, Any
|
3 |
import itertools
|
4 |
from nltk import sent_tokenize
|
5 |
+
# import torch
|
6 |
import nltk
|
|
|
7 |
|
8 |
class PreTrainedPipeline():
|
9 |
|
11 |
# IMPLEMENT_THIS
|
12 |
# Preload all the elements you are going to need at inference.
|
13 |
# For instance your model, processors, tokenizer that might be needed.
|
14 |
+
# This function is only called once, so do all the heavy processing I/O here"""
|
15 |
nltk.download('punkt')
|
16 |
+
self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
|
17 |
self.tokenizer = AutoTokenizer.from_pretrained(path)
|
18 |
|
19 |
self.model_type="t5"
|
20 |
+
# self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
21 |
+
self.device = "cpu"
|
22 |
+
|
23 |
+
self.model.to(self.device)
|
24 |
+
|
25 |
|
26 |
|
27 |
+
def __call__(self, inputs: str, max_words_per_answer: int = 3):
|
28 |
if len(inputs) == 0: return []
|
29 |
inputs = " ".join(inputs.split())
|
30 |
sents, answers = self._extract_answers(inputs)
|
33 |
if len(flat_answers) == 0:
|
34 |
return []
|
35 |
|
36 |
+
questions, qg_examples = self.prepare_and_generate_questions(sents, answers)
|
37 |
+
output = [{'answer': example['answer'], 'question': que} for example, que in zip(qg_examples, questions)]
|
38 |
+
output = self.clean_generated_QAs(output, max_words_per_answer)
|
39 |
+
return output
|
40 |
+
|
41 |
+
def prepare_and_generate_questions(self, sents, answers):
|
42 |
qg_examples = self._prepare_inputs_for_qg_from_answers_hl(sents, answers)
|
43 |
+
|
44 |
qg_inputs = [example['source_text'] for example in qg_examples]
|
45 |
questions = self._generate_questions(qg_inputs)
|
46 |
+
return questions, qg_examples
|
47 |
+
|
48 |
+
|
49 |
+
def clean_answers_list_of_lists(self, answers):
|
50 |
+
clean_answers = []
|
51 |
+
for answer_list in answers:
|
52 |
+
answer_list = answer_list[:-1]
|
53 |
+
answer_list = list(set([a.strip() for a in answer_list]))
|
54 |
+
clean_answers.append(answer_list)
|
55 |
+
return clean_answers
|
56 |
+
|
57 |
|
58 |
def _extract_answers(self, context):
|
|
|
59 |
sents, inputs = self._prepare_inputs_for_ans_extraction(context)
|
60 |
inputs = self._tokenize(inputs, padding=True, truncation=True)
|
61 |
|
67 |
|
68 |
dec = [self.tokenizer.decode(ids, skip_special_tokens=False) for ids in outs]
|
69 |
answers = [item.split('<sep>') for item in dec]
|
70 |
+
|
71 |
+
answers = self.clean_answers_list_of_lists(answers)
|
72 |
+
|
73 |
return sents, answers
|
74 |
|
75 |
+
|
76 |
|
77 |
def _prepare_inputs_for_ans_extraction(self, text):
|
|
|
78 |
sents = sent_tokenize(text)
|
79 |
|
80 |
inputs = []
|
111 |
return inputs
|
112 |
|
113 |
def _generate_questions(self, inputs):
|
|
|
114 |
inputs = self._tokenize(inputs, padding=True, truncation=True)
|
115 |
|
116 |
outs = self.model.generate(
|
122 |
|
123 |
questions = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in outs]
|
124 |
return questions
|
|
|
|
|
125 |
|
126 |
def _prepare_inputs_for_qg_from_answers_hl(self, sents, answers):
|
|
|
127 |
inputs = []
|
128 |
for i, answer in enumerate(answers):
|
129 |
if len(answer) == 0: continue
|
132 |
sents_copy = sents[:]
|
133 |
answer_text = self.remove_pad(answer_text)
|
134 |
answer_text = answer_text.strip()
|
|
|
|
|
135 |
|
136 |
try:
|
137 |
ans_start_idx = sent.lower().index(answer_text.lower())
|
151 |
|
152 |
return inputs
|
153 |
|
154 |
+
def clean_generated_QAs(self, generated_QAs, max_words_per_answer):
|
155 |
clean_QAs = []
|
156 |
answers_used = set()
|
157 |
# Only allow 1 question per answer, take the first case of it
|
158 |
for qa in generated_QAs:
|
159 |
+
answer_word_length = len(qa['answer'].strip().split())
|
160 |
+
if qa['answer'] in answers_used or answer_word_length > max_words_per_answer:
|
161 |
+
continue
|
162 |
answers_used.add(qa['answer'])
|
163 |
clean_QAs.append(qa)
|
164 |
return clean_QAs
|
166 |
def remove_pad(self, str):
|
167 |
if "<pad>" in str:
|
168 |
return str.replace("<pad>", "")
|
169 |
+
return str
|
|