File size: 6,908 Bytes
cd6775d | 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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | from retrieval import reader_model
import pathlib, os
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
device = "cuda"
import torch
import regex
import string
from sentence_transformers import CrossEncoder
import requests
model_cross_encoder = CrossEncoder('cross-encoder/quora-roberta-base',device=device)
model_cross_encoder.model.eval()
def normalize_answer(s):
def remove_articles(text):
return regex.sub(r'\b(a|an|the)\b', ' ', text)
def white_space_fix(text):
return ' '.join(text.split())
def remove_punc(text):
exclude = set(string.punctuation)
return ''.join(ch for ch in text if ch not in exclude)
def lower(text):
return text.lower()
return white_space_fix(remove_articles(remove_punc(lower(s))))
def match_or_not(prediction, ground_truth):
norm_predict = normalize_answer(prediction)
norm_answer = normalize_answer(ground_truth)
return norm_answer in norm_predict
def have_seen_or_not(query_item,query_seen_list,query_type):
if 'Unsolved' in query_type:
return False
for query_seen in query_seen_list:
if model_cross_encoder.predict([(query_seen, query_item)]) > 0.5:
return True
return False
if __name__ == '__main__':
import socket
print('Loading data....')
HOST = '127.0.0.1'
PORT = 50002
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind((HOST, PORT))
sock.listen(5)
print('Waiting for connection...')
sum_cite = 0
good_cite = 0
dic_question_answer_to_reference = []
ques_idx = 0
#start_idx = 0
with torch.no_grad():
while True:
connection,address = sock.accept()
print('connect success from {}'.format(address))
continue_label = True
query_seen_list = []
start = True
break_flag = False
while continue_label:
continue_label = False
#try:
#connection.settimeout(5)
buf = connection.recv(10240)
query = buf.decode()
print('recv query is {}'.format(query))
if query == 'end':
break_flag = True
break
query_list = query.split('\n')
message = ''
for idx in range(len(query_list)):
query_item = query_list[idx]
if 'Query' in query_item and ']:' in query_item:
temp = query_item.split(']')
if len(temp) < 2:
continue
query_type = temp[0]
query_item = temp[1]
if ':' in query_item:
query_item = query_item[1:]
print('solving: '+query_item)
if not have_seen_or_not(query_item,query_seen_list,query_type):
now_reference = {}
query_seen_list.append(query_item)
url = 'http://localhost:8894/api/search?query='+query_item+'&k=1'
response = requests.get(url=url)
res_dic = response.json()
# corpus_list_topk = res_dic['topk']
# #print(corpus_list_topk)
top1_passage = res_dic['text']
# print(top1_passage)
#top1_passage = retrieval_model_hotpotqa.rerank_topk_colbert(corpus_list_topk, query_item)
answer,relevance_score = reader_model.get_answer(query=query_item,texts='',title=top1_passage)
now_reference['query'] = query_item
now_reference['answer'] = answer
now_reference['reference'] = top1_passage
now_reference['ref_score'] = relevance_score
now_reference['idx'] = ques_idx
dic_question_answer_to_reference.append(now_reference)
print('answer is '+answer)
print('reference is'+top1_passage)
print('score is {}'.format(relevance_score))
sum_cite += 1
print('query_type is '+query_type)
if 'Unsolved' in query_type:
message = '[Unsolved Query]:{}<SEP>[Answer]:{}<SEP>[Reference]:{}<SEP>'.format(query_item,
answer,
top1_passage)
print(message)
continue_label = True
if relevance_score > 1.5:
good_cite += 1
break
elif relevance_score > 1.5:
good_cite += 1
answer_start_idx = idx+1
predict_answer = ''
while answer_start_idx < len(query_list):
if 'Answer' in query_list[answer_start_idx]:
predict_answer = query_list[answer_start_idx]
break
answer_start_idx += 1
print('predict answer is '+predict_answer)
match_label = match_or_not(prediction=predict_answer,ground_truth=answer)
if match_label:
continue
else:
message = '[Query]:{}<SEP>[Answer]:{}<SEP>[Reference]:{}<SEP>'.format(query_item,
answer,
top1_passage)
print(message)
continue_label = True
break
if continue_label:
connection.send(message.encode())
else:
connection.send('end'.encode())
while True:
data = connection.recv(1024)
if not data:
break
if not break_flag:
ques_idx += 1
connection.close() |