def main(args):
saved_state = load_states_from_checkpoint(args.model_file)
set_encoder_params_from_state(saved_state.encoder_params, args)
tensorizer, encoder, _ = init_biencoder_components(args.encoder_model_type,
args,
inference_only=True)
encoder = encoder.question_model
encoder, _ = setup_for_distributed_mode(encoder, None, args.device,
args.n_gpu, args.local_rank,
args.fp16)
encoder.eval()
# load weights from the model file
model_to_load = get_model_obj(encoder)
logger.info('Loading saved model state ...')
prefix_len = len('question_model.')
question_encoder_state = {
key[prefix_len:]: value
for (key, value) in saved_state.model_dict.items()
if key.startswith('question_model.')
}
model_to_load.load_state_dict(question_encoder_state)
vector_size = model_to_load.get_out_size()
logger.info('Encoder vector_size=%d', vector_size)
if args.hnsw_index:
index = DenseHNSWFlatIndexer(vector_size, args.index_buffer)
else:
index = DenseFlatIndexer(vector_size, args.index_buffer)
retriever = DenseRetriever(encoder, args.batch_size, tensorizer, index)
# index all passages
ctx_files_pattern = args.encoded_ctx_file
input_paths = glob.glob(ctx_files_pattern)
index_path = "_".join(input_paths[0].split("_")[:-1])
if args.save_or_load_index and (os.path.exists(index_path) or
os.path.exists(index_path + ".index.dpr")):
retriever.index.deserialize_from(index_path)
else:
logger.info('Reading all passages data from files: %s', input_paths)
retriever.index.index_data(input_paths)
if args.save_or_load_index:
retriever.index.serialize(index_path)
# get questions & answers
questions = []
question_ids = []
question_answers = []
for ds_item in parse_qa_csv_file(args.qa_file):
question_id, question, answers = ds_item
question_ids.append(question_id)
questions.append(question)
question_answers.append(answers)
questions_tensor = retriever.generate_question_vectors(questions)
# get top k results
top_ids_and_scores = retriever.get_top_docs(questions_tensor.numpy(),
args.n_docs)
all_passages = load_passages(args.ctx_file)
if len(all_passages) == 0:
raise RuntimeError(
'No passages data found. Please specify ctx_file param properly.')
questions_doc_hits = validate(all_passages, question_answers,
top_ids_and_scores, args.validation_workers,
args.match)
if args.out_file:
save_results(all_passages, questions, question_ids, question_answers,
top_ids_and_scores, questions_doc_hits, args.out_file)
if args.retrieval_type == "tfidf":
import drqa_retriever as retriever
ranker = retriever.get_class('tfidf')(tfidf_path=args.tfidf_path)
top_ids_and_scores = []
for question in questions:
psg_ids, scores = ranker.closest_docs(question, args.n_docs)
top_ids_and_scores.append((psg_ids, scores))
else:
from dpr.models import init_biencoder_components
from dpr.utils.data_utils import Tensorizer
from dpr.utils.model_utils import setup_for_distributed_mode, get_model_obj, load_states_from_checkpoint
from dpr.indexer.faiss_indexers import DenseIndexer, DenseHNSWFlatIndexer, DenseFlatIndexer
from dense_retriever import DenseRetriever
saved_state = load_states_from_checkpoint(args.dpr_model_file)
set_encoder_params_from_state(saved_state.encoder_params, args)
tensorizer, encoder, _ = init_biencoder_components(
args.encoder_model_type, args, inference_only=True)
encoder = encoder.question_model
setup_args_gpu(args)
encoder, _ = setup_for_distributed_mode(encoder, None, args.device,
args.n_gpu, args.local_rank,
args.fp16)
encoder.eval()
# load weights from the model file
model_to_load = get_model_obj(encoder)
prefix_len = len('question_model.')
question_encoder_state = {
key[prefix_len:]: value
for (key, value) in saved_state.model_dict.items()
def main(args):
saved_state = load_states_from_checkpoint(args.model_file)
set_encoder_params_from_state(saved_state.encoder_params, args)
print_args(args)
tensorizer, encoder, _ = init_biencoder_components(
args.encoder_model_type, args, inference_only=True
)
encoder = encoder.ctx_model
encoder, _ = setup_for_distributed_mode(
encoder,
None,
args.device,
args.n_gpu,
args.local_rank,
args.fp16,
args.fp16_opt_level,
)
encoder.eval()
# load weights from the model file
model_to_load = get_model_obj(encoder)
logger.info("Loading saved model state ...")
logger.debug("saved model keys =%s", saved_state.model_dict.keys())
prefix_len = len("ctx_model.")
ctx_state = {
key[prefix_len:]: value
for (key, value) in saved_state.model_dict.items()
if key.startswith("ctx_model.")
}
model_to_load.load_state_dict(ctx_state)
logger.info("reading data from file=%s", args.ctx_file)
rows = []
csv.field_size_limit(sys.maxsize)
with open(args.ctx_file) as tsvfile:
reader = csv.reader(tsvfile, delimiter="\t")
# file format: doc_id, doc_text, title
rows.extend([(row[0], row[1], row[2]) for row in reader if row[0] != "id"])
shard_size = int(len(rows) / args.num_shards)
start_idx = args.shard_id * shard_size
end_idx = start_idx + shard_size
logger.info(
"Producing encodings for passages range: %d to %d (out of total %d)",
start_idx,
end_idx,
len(rows),
)
rows = rows[start_idx:end_idx]
data = gen_ctx_vectors(rows, encoder, tensorizer, True)
file = args.out_file + "_" + str(args.shard_id) + ".pkl"
pathlib.Path(os.path.dirname(file)).mkdir(parents=True, exist_ok=True)
logger.info("Writing results to %s" % file)
with open(file, mode="wb") as f:
pickle.dump(data, f)
logger.info("Total passages processed %d. Written to %s", len(data), file)
def main(args):
if not args.encoder_model_type == 'hf_attention' and not args.encoder_model_type == 'colbert':
saved_state = load_states_from_checkpoint(args.model_file)
set_encoder_params_from_state(saved_state.encoder_params, args)
print_args(args)
tensorizer, encoder, _ = init_biencoder_components(args.encoder_model_type,
args,
inference_only=True)
if not args.encoder_model_type == 'hf_attention' and not args.encoder_model_type == 'colbert':
encoder = encoder.ctx_model
encoder, _ = setup_for_distributed_mode(encoder, None, args.device,
args.n_gpu, args.local_rank,
args.fp16, args.fp16_opt_level)
encoder.eval()
# load weights from the model file
model_to_load = get_model_obj(encoder)
logger.info('Loading saved model state ...')
#logger.debug('saved model keys =%s', saved_state.model_dict.keys())
if not args.encoder_model_type == 'hf_attention' and not args.encoder_model_type == 'colbert':
prefix_len = len('ctx_model.')
ctx_state = {
key[prefix_len:]: value
for (key, value) in saved_state.model_dict.items()
if key.startswith('ctx_model.')
}
model_to_load.load_state_dict(ctx_state)
logger.info('reading data from file=%s', args.ctx_file)
rows = []
with open(args.ctx_file) as tsvfile:
reader = csv.reader(tsvfile, delimiter='\t')
for k, row in enumerate(reader):
#if k == 1000:
# break
if not row[0] == 'id':
rows.append((row[0], row[1], row[2]))
#rows.extend([(row[0], row[1], row[2]) for row in reader if row[0] != 'id'])
shard_size = int(len(rows) / args.num_shards)
start_idx = args.shard_id * shard_size
end_idx = start_idx + shard_size
logger.info(
'Producing encodings for passages range: %d to %d (out of total %d)',
start_idx, end_idx, len(rows))
rows = rows[start_idx:end_idx]
data = gen_ctx_vectors(rows, encoder, tensorizer, True)
file = args.out_file + '_' + str(args.shard_id)
pathlib.Path(os.path.dirname(file)).mkdir(parents=True, exist_ok=True)
logger.info('Writing results to %s' % file)
with open(file, mode='wb') as f:
pickle.dump(data, f)
logger.info('Total passages processed %d. Written to %s', len(data), file)