preprocess_reader_data.py

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
 Reader data preprocessor command line tool
"""
import argparse
import logging

from dpr.data.reader_data import convert_retriever_results
from dpr.models import init_tenzorizer
from dpr.options import print_args, add_encoder_params, add_reader_preprocessing_params, add_tokenizer_params

logger = logging.getLogger()


def main(args):
    tensorizer = init_tenzorizer(args.encoder_model_type, args)

    # disable auto-padding to save disk space of serialized files
    tensorizer.set_pad_to_max(False)

    convert_retriever_results(args.is_train_set, args.retriever_results, args.out_file, args.gold_passages_src,
                              tensorizer, args.num_workers)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    add_encoder_params(parser)
    add_tokenizer_params(parser)
    add_reader_preprocessing_params(parser)

    parser.add_argument("--is_train_set", action='store_true',
                        help="If true, the data will be binarised for train model usage (split into ctx+ and ctx- \
                        and with answer spans selected)")
    parser.add_argument("--retriever_results", required=True, type=str,
                        help="File with retriever results file(json format)")
    parser.add_argument("--out_file", required=True, type=str, help="The file to write serialized results to")

    args = parser.parse_args()

    print_args(args)

    main(args)