Spaces:
Runtime error
Runtime error
File size: 4,600 Bytes
9f23e0b |
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 |
import torch
import numpy as np
# # from transformers import AutoTokenizer, AutoModelForQuestionAnswering
class QuestionAnswer:
def __init__(self, data, model, tokenizer, torch_device):
self.max_length = 384
self.doc_stride = 128
self.tokenizer = tokenizer
self.model = model
self.data = data
self.torch_device = torch_device
self.output = None
self.features = None
self.results = None
def get_output_from_model(self):
# data = {'question': question, 'context': context}
with torch.no_grad():
tokenized_data = self.tokenizer(
self.data['question'],
self.data['context'],
truncation='only_second',
max_length=self.max_length,
stride=self.doc_stride,
return_overflowing_tokens=True,
return_offsets_mapping=True,
padding='max_length',
return_tensors='pt'
).to(self.torch_device)
output = self.model(tokenized_data['input_ids'], tokenized_data['attention_mask'])
return output
# print(output.keys())
# print(output['start_logits'].shape)
# print(output['end_logits'].shape)
# print(tokenized_data.keys())
def prepare_features(self, example):
tokenized_example = self.tokenizer(
example['question'],
example['context'],
truncation='only_second',
max_length=self.max_length,
stride=self.doc_stride,
return_overflowing_tokens=True,
return_offsets_mapping=True,
padding='max_length',
)
# sample_mapping = tokenized_example.pop("overflow_to_sample_mapping")
for i in range(len(tokenized_example['input_ids'])):
sequence_ids = tokenized_example.sequence_ids(i)
# print(sequence_ids)
context_index = 1
# sample_index = sample_mapping[i]
tokenized_example["offset_mapping"][i] = [
(o if sequence_ids[k] == context_index else None)
for k, o in enumerate(tokenized_example["offset_mapping"][i])
]
return tokenized_example
def postprocess_qa_predictions(self, data, features, raw_predictions, top_n_answers=5, max_answer_length=30):
all_start_logits, all_end_logits = raw_predictions.start_logits, raw_predictions.end_logits
# print(all_start_logits)
results = []
context = data['context']
# print(len(features['input_ids']))
for i in range(len(features['input_ids'])):
start_logits = all_start_logits[i].cpu().numpy()
end_logits = all_end_logits[i].cpu().numpy()
# print(start_logits)
offset_mapping = features['offset_mapping'][i]
start_indices = np.argsort(start_logits)[-1: -top_n_answers - 1: -1].tolist()
end_indices = np.argsort(end_logits)[-1: -top_n_answers - 1: -1].tolist()
for start_index in start_indices:
for end_index in end_indices:
if (
start_index >= len(offset_mapping)
or end_index >= len(offset_mapping)
or offset_mapping[start_index] is None
or offset_mapping[end_index] is None
or end_index < start_index
or end_index - start_index + 1 > max_answer_length
):
continue
start_char = offset_mapping[start_index][0]
end_char = offset_mapping[end_index][1]
# print(start_logits[start_index])
# print(end_logits[end_index])
score = start_logits[start_index] + end_logits[end_index]
results.append(
{
'score': float('%.*g' % (3, score)),
'text': context[start_char: end_char]
}
)
results = sorted(results, key=lambda x: x["score"], reverse=True)[:top_n_answers]
return results
def get_results(self):
self.output = self.get_output_from_model()
self.features = self.prepare_features(self.data)
self.results = self.postprocess_qa_predictions(self.data, self.features, self.output)
return self.results |