convert_pyserini_to_dpr_file.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.

"""
 Command line tool to get dense results and validate them
"""

import argparse
import os
import csv
import glob
import json
import gzip
import logging
import pickle
import time
from typing import List, Tuple, Dict, Iterator
from tqdm import tqdm
from pyserini.index import IndexReader

import numpy as np
import torch
from torch import Tensor as T
from torch import nn

from dpr.data.qa_validation import calculate_matches
from dpr.models import init_biencoder_components
from dpr.options import (
    add_encoder_params,
    setup_args_gpu,
    print_args,
    set_encoder_params_from_state,
    add_tokenizer_params,
    add_cuda_params,
)
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,
    DenseReconIndexer,
)

logger = logging.getLogger()
logger.setLevel(logging.INFO)
if logger.hasHandlers():
    logger.handlers.clear()
console = logging.StreamHandler()
logger.addHandler(console)

def load_passages(ctx_file: str) -> Dict[object, Tuple[str, str]]:
    docs = {}
    logger.info("Reading data from: %s", ctx_file)
    if ctx_file.endswith(".gz"):
        with gzip.open(ctx_file, "rt") as tsvfile:
            reader = csv.reader(
                tsvfile,
                delimiter="\t",
            )
            # file format: doc_id, doc_text, title
            for row in reader:
                if row[0] != "id":
                    docs[row[0]] = (row[1], row[2])
    else:
        with open(ctx_file) as tsvfile:
            reader = csv.reader(
                tsvfile,
                delimiter="\t",
            )
            # file format: doc_id, doc_text, title
            for row in reader:
                if row[0] != "id":
                    docs[row[0]] = (row[1], row[2])
    return docs

def validate(
    passages: Dict[object, Tuple[str, str]],
    answers: List[List[str]],
    result_ctx_ids: List[Tuple[List[object], List[float]]],
    workers_num: int,
    match_type: str,
) -> List[List[bool]]:
    match_stats = calculate_matches(
        passages, answers, result_ctx_ids, workers_num, match_type
    )
    top_k_hits = match_stats.top_k_hits

    logger.info("Validation results: top k documents hits %s", top_k_hits)
    top_k_hits = [v / len(result_ctx_ids) for v in top_k_hits]
    logger.info("Validation results: top k documents hits accuracy %s", top_k_hits)
    return match_stats.questions_doc_hits

def save_results(
    path,
    passages: Dict[object, Tuple[str, str]],
    questions: List[str],
    answers: List[List[str]],
    top_passages_and_scores: List[Tuple[List[object], List[float]]],
    per_question_hits: List[List[bool]],
    out_file: str,
):
    # join passages text with the result ids, their questions and assigning has|no answer labels
    merged_data = []
    assert len(per_question_hits) == len(questions) == len(answers)
    all_keys = list(passages.keys())
    for i, q in tqdm(enumerate(questions)):
        q_answers = answers[i]
        results_and_scores = top_passages_and_scores[i]
        hits = per_question_hits[i]
        docs = [passages[doc_id] for doc_id in results_and_scores[0]]
        scores = [str(score) for score in results_and_scores[1]]
        ctxs_num = len(hits)
        
        positive_ctxs = []
        negative_ctxs = []
        hard_negative_ctxs = []
        for c in range(ctxs_num):
          sample = {"title": docs[c][1],
                    "title_score": 0,
                    "text": docs[c][0],
                    "score": scores[c],
                    "passage_id": results_and_scores[0][c]
                    }
          if hits[c]:
            positive_ctxs.append(sample) 
          else:
            hard_negative_ctxs.append(sample)
        
        if len(positive_ctxs) == 0:
          continue
        
        docids = {doc_id:0 for doc_id in results_and_scores[0]}
        while len(negative_ctxs) < 50:
          neg_id = all_keys[np.random.randint(len(all_keys))]
          if neg_id in docids:
            continue
          else:
            negative_ctxs.append({"title": passages[neg_id][1],
                                  "title_score": 0,
                                  "text": passages[neg_id][0],
                                  "score": 0,
                                  "passage_id": neg_id
                                  })
        merged_data.append(
            {
                "dataset": path,
                "question": q,
                "answers": q_answers,
                "positive_ctxs": positive_ctxs,
                "negative_ctxs": negative_ctxs,
                "hard_negative_ctxs": hard_negative_ctxs,
            }
        )
    
    print(f"Filtered data:{len(merged_data)}/{len(questions)}")

    with open(out_file, "w") as writer:
        writer.write(json.dumps(merged_data, indent=4) + "\n")
    logger.info("Saved results * scores  to %s", out_file)

def main(args):
    logger.info("Loading retrieved results ...")
    with open(args.retrieval_results) as f:
        retrieved_data = json.load(f)
    
    questions = []
    question_answers = []
    top_ids_and_scores = []
    
    for sample in retrieved_data.values():
      questions.append(sample["question"])
      question_answers.append(sample["answers"])
      scores, ids = [], []
      for ctx in sample["contexts"]:
        scores.append(ctx["score"])
        ids.append(ctx["docid"])
      top_ids_and_scores.append((ids, scores))

    logger.info("Loading passages ...")
    all_passages = load_passages(args.ctx_file)
    
    questions_doc_hits = validate(
        all_passages,
        question_answers,
        top_ids_and_scores,
        args.validation_workers,
        args.match,
    )

    if args.out_file:
        save_results(
            args.retrieval_results,
            all_passages,
            questions,
            question_answers,
            top_ids_and_scores,
            questions_doc_hits,
            args.out_file,
        )


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

    add_encoder_params(parser)
    add_tokenizer_params(parser)
    add_cuda_params(parser)
    
    parser.add_argument(
        "--ctx_file",
        required=True,
        type=str,
        default=None,
        help="All passages file in the tsv format: id \\t passage_text \\t title",
    )

    parser.add_argument(
        "--out_file",
        type=str,
        default=None,
        help="output .tsv file path to write results to ",
    )
    parser.add_argument(
        "--match",
        type=str,
        default="string",
        choices=["regex", "string"],
        help="Answer matching logic type",
    )
    parser.add_argument(
        "--validation_workers",
        type=int,
        default=16,
        help="Number of parallel processes to validate results",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=32,
        help="Batch size for question encoder forward pass",
    )
    parser.add_argument(
        "--retrieval_results",
        required=True,
        type=str,
        default=None,
        help="Retreival results from another model",
    )

    args = parser.parse_args()

    setup_args_gpu(args)
    print_args(args)
    main(args)