Function compare_server_and_bruteforce

tests/test_multiple_datasets.py:255–300

(bf_csv_path: str)

def compare_server_and_bruteforce(bf_csv_path: str):
    bf = pl.read_csv(bf_csv_path)
    base_dir, bf_file_name = os.path.split(bf_csv_path)
    server_file_name = bf_file_name.replace("bf_", "server_")
    server_path = os.path.join(base_dir, server_file_name)
    sr = pl.read_csv(server_path)

    total_queries = sr.shape[0]
    hits = 0
    recall_sum = 0.0

    for row in sr.iter_rows(named=True):
        qid = row["query_id"]
        match = bf.filter(pl.col("query_id") == qid).row(0)

        bf_top_ids = [match[1], match[3], match[5], match[7]]
        bf_top_sims = [match[2], match[4], match[6], match[8]]
        srv_top_ids = [row["top1_id"], row["top2_id"], row["top3_id"], row["top4_id"]]
        srv_top_sims = [
            row["top1_sim"],
            row["top2_sim"],
            row["top3_sim"],
            row["top4_sim"],
        ]

        print(colored(f"Query ID: {qid}", "yellow"))
        print(
            colored(f"BF top4 IDs: {bf_top_ids}", "blue"),
            colored(f"sims: {bf_top_sims}", "blue"),
        )
        print(
            colored(f"Server top4 IDs: {srv_top_ids}", "magenta"),
            colored(f"sims: {srv_top_sims}", "magenta"),
        )
        print("-" * 40)

        top_srv = set(srv_top_ids)
        top_bf = set(bf_top_ids)
        common = top_srv.intersection(top_bf)
        hits += len(common)
        recall_sum += len(common) / 4

    print(colored(f"Total queries compared: {total_queries}", "cyan"))
    print(colored(f"Total matching embeddings in server top4: {hits}", "cyan"))
    recall_percentage = (recall_sum / total_queries) * 100
    print(colored(f"Mean recall: {recall_percentage:.2f}%", "green"))


if __name__ == "__main__":