INNER CODE UNIT · Python

save_qrels

UKPLab/gpl · gpl/toolkit/beir.py:28

def save_qrels(qrels: Dict[str, Dict[str, float]], output_dir, split):
    os.makedirs(os.path.join(output_dir, "qrels"), exist_ok=True)

    assert split in ["train", "test", "dev"]
    save_path = os.path.join(output_dir, f"qrels/{split}.tsv")
    with open(save_path, "w") as f:
        header = "query-id\tcorpus-id\tscore\n"
        f.write(header)
        for qid, rels in qrels.items():
            for doc_id, score in rels.items():
                line = f"{qid}\t{doc_id}\t{score}\n"
                f.write(line)

    logger.info(f"Saved (copied) qrels into {save_path}")

View source record →