Spaces:
Running
Running
File size: 2,647 Bytes
9bf1d31 |
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 |
from ..base_handler import ModelHandler
from transformers import AutoTokenizer
import torch
import time
class QuestionAnsweringHandler(ModelHandler):
def __init__(self, model_name, model_class, quantization_type, test_text):
super().__init__(model_name, model_class, quantization_type, test_text)
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
def run_inference(self, model, text):
parts = text.split('QUES')
context = parts[0].strip()
question = parts[1].strip()
inputs = self.tokenizer(question, context, return_tensors='pt', truncation=True, padding=True).to(self.device)
start_time = time.time()
with torch.no_grad():
outputs = model(**inputs)
end_time = time.time()
return outputs, end_time - start_time
def decode_output(self, outputs):
start_logits = outputs.start_logits
end_logits = outputs.end_logits
answer_start = torch.argmax(start_logits)
answer_end = torch.argmax(end_logits) + 1
input_ids = self.tokenizer.encode(self.test_text)
answer = self.tokenizer.decode(input_ids[answer_start:answer_end])
return f"Answer: {answer}"
def compare_outputs(self, original_outputs, quantized_outputs):
"""Compare outputs for question answering models"""
if original_outputs is None or quantized_outputs is None:
return None
orig_start = original_outputs.start_logits.cpu().numpy()
orig_end = original_outputs.end_logits.cpu().numpy()
quant_start = quantized_outputs.start_logits.cpu().numpy()
quant_end = quantized_outputs.end_logits.cpu().numpy()
orig_start_pos = orig_start.argmax()
orig_end_pos = orig_end.argmax()
quant_start_pos = quant_start.argmax()
quant_end_pos = quant_end.argmax()
input_ids = self.tokenizer.encode(self.test_text)
original_answer = self.tokenizer.decode(input_ids[orig_start_pos:orig_end_pos + 1])
quantized_answer = self.tokenizer.decode(input_ids[quant_start_pos:quant_end_pos + 1])
metrics = {
'original_answer': original_answer,
'quantized_answer': quantized_answer,
'start_position_match': float(orig_start_pos == quant_start_pos),
'end_position_match': float(orig_end_pos == quant_end_pos),
'start_logits_mse': ((orig_start - quant_start) ** 2).mean(),
'end_logits_mse': ((orig_end - quant_end) ** 2).mean(),
}
return metrics |