import json
import os
import re
from collections import defaultdict
from multiprocessing import Pool, cpu_count
from time import time
import faiss
import numpy as np
import pandas as pd
import tensorflow as tf
from tqdm import tqdm
import gpt2_estimator
from docproduct.dataset import convert_text_to_feature
from docproduct.models import MedicalQAModelwithBert
from docproduct.tokenization import FullTokenizer
from keras_bert.loader import checkpoint_loader
def load_weight(model, bert_ffn_weight_file=None, ffn_weight_file=None):
if bert_ffn_weight_file:
model.load_weights(bert_ffn_weight_file)
elif ffn_weight_file:
loader = checkpoint_loader(ffn_weight_file)
model.get_layer('q_ffn').set_weights(
[loader('q_ffn/ffn_layer/kernel/.ATTRIBUTES/VARIABLE_VALUE'),
loader('q_ffn/ffn_layer/bias/.ATTRIBUTES/VARIABLE_VALUE')])
model.get_layer('a_ffn').set_weights(
[loader('a_ffn/ffn_layer/kernel/.ATTRIBUTES/VARIABLE_VALUE'),
loader('a_ffn/ffn_layer/bias/.ATTRIBUTES/VARIABLE_VALUE')]
)
class QAEmbed(object):
def __init__(
self,
hidden_size=768,
dropout=0.2,
residual=True,
pretrained_path=None,
batch_size=128,
max_seq_length=256,
ffn_weight_file=None,
bert_ffn_weight_file=None,
load_pretrain=True,
with_question=True,
with_answer=True):
super(QAEmbed, self).__init__()
config_file = os.path.join(pretrained_path, 'bert_config.json')
if load_pretrain:
checkpoint_file = os.path.join(
pretrained_path, 'biobert_model.ckpt')
else:
checkpoint_file = None
# the ffn model takes 2nd to last layer
if bert_ffn_weight_file is None:
layer_ind = -2
else:
layer_ind = -1
self.model = MedicalQAModelwithBert(
hidden_size=768,
dropout=0.2,
residual=True,
config_file=config_file,
checkpoint_file=checkpoint_file,
layer_ind=layer_ind)
self.batch_size = batch_size
self.tokenizer = FullTokenizer(
os.path.join(pretrained_path, 'vocab.txt'))
self.max_seq_length = max_seq_length
# build mode in order to load
question = 'fake' if with_question else None
answer = 'fake' if with_answer else None
self.predict(questions=question, answers=answer, dataset=False)
load_weight(self.model, bert_ffn_weight_file, ffn_weight_file)
def _type_check(self, inputs):
if inputs is not None:
if isinstance(inputs, str):
inputs = [inputs]
elif isinstance(inputs, list):
pass
else:
raise TypeError(
'inputs are supposed to be str of list of str, got {0} instead.'.format(type(inputs)))
return inputs
def _make_inputs(self, questions=None, answers=None, dataset=True):
if questions:
data_size = len(questions)
q_feature_dict = defaultdict(list)
for q in questions:
q_feature = convert_text_to_feature(
q, tokenizer=self.tokenizer, max_seq_length=self.max_seq_length)
q_feature_dict['q_input_ids'].append(q_feature[0])
q_feature_dict['q_input_masks'].append(q_feature[1])
q_feature_dict['q_segment_ids'].append(q_feature[2])
if answers:
data_size = len(answers)
a_feature_dict = defaultdict(list)
for a in answers:
a_feature = convert_text_to_feature(
a, tokenizer=self.tokenizer, max_seq_length=self.max_seq_length)
a_feature_dict['a_input_ids'].append(a_feature[0])
a_feature_dict['a_input_masks'].append(a_feature[1])
a_feature_dict['a_segment_ids'].append(a_feature[2])
if questions and answers:
q_feature_dict.update(a_feature_dict)
model_inputs = q_feature_dict
elif questions:
model_inputs = q_feature_dict
elif answers:
model_inputs = a_feature_dict
model_inputs = {k: tf.convert_to_tensor(
np.stack(v, axis=0)) for k, v in model_inputs.items()}
if dataset:
model_inputs = tf.data.Dataset.from_tensor_slices(model_inputs)
model_inputs = model_inputs.batch(self.batch_size)
return model_inputs
def predict(self, questions=None, answers=None, dataset=True):
# type check
questions = self._type_check(questions)
answers = self._type_check(answers)
if questions is not None and answers is not None:
assert len(questions) == len(answers)
model_inputs = self._make_inputs(questions, answers, dataset)
model_outputs = []
if dataset:
for batch in tqdm(iter(model_inputs), total=int(len(questions) / self.batch_size)):
model_outputs.append(self.model(batch))
model_outputs = np.concatenate(model_outputs, axis=0)
else:
model_outputs = self.model(model_inputs)
return model_outputs
class FaissTopK(object):
def __init__(self, embedding_file):
super(FaissTopK, self).__init__()
self.embedding_file = embedding_file
_, ext = os.path.splitext(self.embedding_file)
if ext == '.pkl':
self.df = pd.read_pickle(self.embedding_file)
else:
self.df = pd.read_parquet(self.embedding_file)
self._get_faiss_index()
# self.df.drop(columns=["Q_FFNN_embeds", "A_FFNN_embeds"], inplace=True)
def _get_faiss_index(self):
# with Pool(cpu_count()) as p:
# question_bert = p.map(eval, self.df["Q_FFNN_embeds"].tolist())
# answer_bert = p.map(eval, self.df["A_FFNN_embeds"].tolist())
question_bert = self.df["Q_FFNN_embeds"].tolist()
self.df.drop(columns=["Q_FFNN_embeds"], inplace=True)
answer_bert = self.df["A_FFNN_embeds"].tolist()
self.df.drop(columns=["A_FFNN_embeds"], inplace=True)
question_bert = np.array(question_bert, dtype='float32')
answer_bert = np.array(answer_bert, dtype='float32')
self.answer_index = faiss.IndexFlatIP(answer_bert.shape[-1])
self.question_index = faiss.IndexFlatIP(question_bert.shape[-1])
self.answer_index.add(answer_bert)
self.question_index.add(question_bert)
del answer_bert, question_bert
def predict(self, q_embedding, search_by='answer', topk=5, answer_only=True):
if search_by == 'answer':
_, index = self.answer_index.search(
q_embedding.astype('float32'), topk)
else:
_, index = self.question_index.search(
q_embedding.astype('float32'), topk)
output_df = self.df.iloc[index[0], :]
if answer_only:
return output_df.answer.tolist()
else:
return (output_df.question.tolist(), output_df.answer.tolist())
class RetreiveQADoc(object):
def __init__(self,
pretrained_path=None,
ffn_weight_file=None,
bert_ffn_weight_file='models/bertffn_crossentropy/bertffn',
embedding_file='qa_embeddings/bertffn_crossentropy.zip'
):
super(RetreiveQADoc, self).__init__()
self.qa_embed = QAEmbed(
pretrained_path=pretrained_path,
ffn_weight_file=ffn_weight_file,
bert_ffn_weight_file=bert_ffn_weight_file
)
self.faiss_topk = FaissTopK(embedding_file)
def predict(self, questions, search_by='answer', topk=5, answer_only=True):
embedding = self.qa_embed.predict(questions=questions)
return self.faiss_topk.predict(embedding, search_by, topk, answer_only)
def getEmbedding(self, questions, search_by='answer', topk=5, answer_only=True):
embedding = self.qa_embed.predict(questions=questions)
return embedding
class GenerateQADoc(object):
def __init__(self,
pretrained_path='models/pubmed_pmc_470k/',
ffn_weight_file=None,
bert_ffn_weight_file='models/bertffn_crossentropy/bertffn',
gpt2_weight_file='models/gpt2',
embedding_file='qa_embeddings/bertffn_crossentropy.zip'
):
super(GenerateQADoc, self).__init__()
tf.compat.v1.disable_eager_execution()
session_config = tf.compat.v1.ConfigProto(
allow_soft_placement=True)
session_config.gpu_options.allow_growth = False
config = tf.estimator.RunConfig(
session_config=session_config)
self.batch_size = 1
self.gpt2_weight_file = gpt2_weight_file
gpt2_model_fn = gpt2_estimator.get_gpt2_model_fn(
accumulate_gradients=5,
learning_rate=0.1,
length=512,
batch_size=self.batch_size,
temperature=0.7,
top_k=0
)
hparams = gpt2_estimator.default_hparams()
with open(os.path.join(gpt2_weight_file, 'hparams.json')) as f:
hparams.override_from_dict(json.load(f))
self.estimator = tf.estimator.Estimator(
gpt2_model_fn,
model_dir=gpt2_weight_file,
params=hparams,
config=config)
self.encoder = gpt2_estimator.encoder.get_encoder(gpt2_weight_file)
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
self.embed_sess = tf.compat.v1.Session(config=config)
with self.embed_sess.as_default():
self.qa_embed = QAEmbed(
pretrained_path=pretrained_path,
ffn_weight_file=ffn_weight_file,
bert_ffn_weight_file=bert_ffn_weight_file,
with_answer=False,
load_pretrain=False
)
self.faiss_topk = FaissTopK(embedding_file)
def _get_gpt2_inputs(self, question, questions, answers):
assert len(questions) == len(answers)
line = '`QUESTION: %s `ANSWER: ' % question
for q, a in zip(questions, answers):
line = '`QUESTION: %s `ANSWER: %s ' % (q, a) + line
return line
def predict(self, questions, search_by='answer', topk=5, answer_only=False):
embedding = self.qa_embed.predict(
questions=questions, dataset=False).eval(session=self.embed_sess)
if answer_only:
topk_answer = self.faiss_topk.predict(
embedding, search_by, topk, answer_only)
else:
topk_question, topk_answer = self.faiss_topk.predict(
embedding, search_by, topk, answer_only)
gpt2_input = self._get_gpt2_inputs(
questions[0], topk_question, topk_answer)
gpt2_pred = self.estimator.predict(
lambda: gpt2_estimator.predict_input_fn(inputs=gpt2_input, batch_size=self.batch_size, checkpoint_path=self.gpt2_weight_file))
raw_output = gpt2_estimator.predictions_parsing(
gpt2_pred, self.encoder)
# result_list = [re.search('`ANSWER:(.*)`QUESTION:', s)
# for s in raw_output]
# result_list = [s for s in result_list if s]
# try:
# r = result_list[0].group(1)
# except (AttributeError, IndexError):
# r = topk_answer[0]
refine1 = re.sub('`QUESTION:.*?`ANSWER:','' , str(raw_output[0]) , flags=re.DOTALL)
refine2 = refine1.split('`QUESTION: ')[0]
return refine2
if __name__ == "__main__":
gen = GenerateQADoc()
print(gen.predict('my eyes hurt'))
