from pathlib import Path from transformers import AutoTokenizer, pipeline import numpy as np import torch import torch.nn as nn from text_utils import post_process_answer from graph_utils import find_best_cluster from optimum.intel import OVModelForQuestionAnswering import os import json from text_utils import * # os.environ['HTTP_PROXY'] = 'http://proxy.hcm.fpt.vn:80' class QAEnsembleModel_modify(nn.Module): # def __init__(self, model_name, model_checkpoints, entity_dict, # thr=0.1, device="cuda:0"): def __init__(self, model_name, entity_dict, thr=0.1, device="cpu"): super(QAEnsembleModel_modify, self).__init__() self.nlps = [] # model_checkpoint = "./data/qa_model_robust.bin" AUTH_TOKEN = "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" # model_checkpoint = "letrunglinh/qa_pnc" model_convert = OVModelForQuestionAnswering.from_pretrained(model_name, export= True, use_auth_token= AUTH_TOKEN) model_convert.half() model_convert.compile() nlp = pipeline('question-answering', model=model_convert, tokenizer=model_name) self.nlps.append(nlp) self.entity_dict = entity_dict self.thr = thr def forward(self, question, texts, ranking_scores=None): if ranking_scores is None: ranking_scores = np.ones((len(texts),)) curr_answers = [] curr_scores = [] best_score = 0 for i, nlp in enumerate(self.nlps): for text, score in zip(texts, ranking_scores): QA_input = { 'question': question, 'context': text } res = nlp(QA_input) print(res) if res["score"] > self.thr: curr_answers.append(res["answer"]) curr_scores.append(res["score"]) res["score"] = res["score"] * score if i == 0: if res["score"] > best_score: answer = res["answer"] best_score = res["score"] if len(curr_answers) == 0: return None curr_answers = [post_process_answer(x, self.entity_dict) for x in curr_answers] answer = post_process_answer(answer, self.entity_dict) new_best_answer = post_process_answer(find_best_cluster(curr_answers, answer), self.entity_dict) return new_best_answer