| 690 | |
| 691 | |
| 692 | def brute_force_sparse( |
| 693 | queries: List[Dict], docs: List[Dict], top_k: int, dataset: str |
| 694 | ) -> List[Dict]: |
| 695 | bf_file = Path("datasets") / f"hybrid_{dataset}" / "sparse_brute_force.pkl" |
| 696 | if bf_file.exists(): |
| 697 | return pickle.loads(bf_file.read_bytes()) |
| 698 | |
| 699 | print("Computing sparse brute-force...") |
| 700 | results = [] |
| 701 | for q in tqdm(queries, desc="Sparse brute-force"): |
| 702 | scores = defaultdict(float) |
| 703 | q_indices_set = set(q["indices"]) |
| 704 | for d in docs: |
| 705 | # Compute dot product between sparse vectors |
| 706 | for i, val in zip(d["indices"], d["values"]): |
| 707 | if i in q_indices_set: |
| 708 | q_idx = q["indices"].index(i) |
| 709 | scores[d["id"]] += val * q["values"][q_idx] |
| 710 | top = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_k] |
| 711 | results.append( |
| 712 | { |
| 713 | "query_id": q["id"], |
| 714 | "top_results": [{"id": i, "score": s} for i, s in top], |
| 715 | } |
| 716 | ) |
| 717 | |
| 718 | bf_file.parent.mkdir(parents=True, exist_ok=True) |
| 719 | bf_file.write_bytes(pickle.dumps(results)) |
| 720 | return results |
| 721 | |
| 722 | def ensure_collection(name: str): |
| 723 | """Create or recreate the collection with proper configuration""" |