INNER CODE UNIT · Python

extract_queries_split

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

def extract_queries_split(queries: Dict[str, str], qrels: Dict[str, Dict[str, float]]):
    qids_to_exclude = set(qrels) - set(queries)
    for qid in qids_to_exclude:
        queries.pop(qid)
    return queries


def save_queries(queries: Dict[str, str], output_dir):
    os.makedirs(output_dir, exist_ok=True)
    save_path = os.path.join(output_dir, "queries.jsonl")
    with open(save_path, "w") as f:
        for qid, query in queries.items():
            line_dict = {"_id": qid, "text": query, "metadata": {}}
            line = json.dumps(line_dict) + "\n"
            f.write(line)

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

View source record →