Build sparse vectors for brute-force comparison
(corpus: Dict, dataset: str)
| 586 | |
| 587 | |
| 588 | def build_sparse_vectors_for_bf(corpus: Dict, dataset: str) -> List[Dict]: |
| 589 | """Build sparse vectors for brute-force comparison""" |
| 590 | cache_file = Path("datasets") / f"hybrid_{dataset}" / "sparse_vectors.pkl" |
| 591 | if cache_file.exists(): |
| 592 | print("Loading cached sparse vectors...") |
| 593 | return pickle.loads(cache_file.read_bytes()) |
| 594 | |
| 595 | print("Tokenising corpus (sparse)...") |
| 596 | stemmer = SnowballStemmer("english") |
| 597 | |
| 598 | def process_doc(doc_item): |
| 599 | doc_id, doc = doc_item |
| 600 | text = f"{doc.get('title', '')} {doc['text']}" |
| 601 | tokens = SimpleTokenizer.tokenize(text) |
| 602 | terms = [ |
| 603 | stemmer.stem_word(t.lower()) |
| 604 | for t in tokens |
| 605 | if t.lower() not in STOPWORDS and t not in PUNCT and len(t) <= 40 |
| 606 | ] |
| 607 | return {"id": doc_id, "tokens": terms, "length": len(terms)} |
| 608 | |
| 609 | docs = [] |
| 610 | with ThreadPoolExecutor(max_workers=MAX_WORKERS) as ex: |
| 611 | futures = [ex.submit(process_doc, item) for item in corpus.items()] |
| 612 | for f in tqdm(as_completed(futures), total=len(futures), desc="Tokenising"): |
| 613 | docs.append(f.result()) |
| 614 | |
| 615 | total_docs = len(docs) |
| 616 | term_doc_freq = defaultdict(int) |
| 617 | total_len = sum(d["length"] for d in docs) |
| 618 | avg_len = total_len / total_docs |
| 619 | for d in docs: |
| 620 | for t in set(d["tokens"]): |
| 621 | term_doc_freq[t] += 1 |
| 622 | |
| 623 | def build_vector(doc): |
| 624 | tf = defaultdict(int) |
| 625 | for tok in doc["tokens"]: |
| 626 | tf[tok] += 1 |
| 627 | indices, values = [], [] |
| 628 | for tok, raw in tf.items(): |
| 629 | idf = compute_bm25_idf(total_docs, term_doc_freq[tok]) |
| 630 | tf_score = compute_bm25_tf(raw, doc["length"], avg_len) |
| 631 | bm25_score = idf * tf_score |
| 632 | if bm25_score > 0: |
| 633 | indices.append(hash(tok) % (2**31)) # Simple hash for index |
| 634 | values.append(bm25_score) |
| 635 | return { |
| 636 | "id": doc["id"], |
| 637 | "text": " ".join(doc["tokens"]), |
| 638 | "indices": indices, |
| 639 | "values": values, |
| 640 | } |
| 641 | |
| 642 | vectors = [] |
| 643 | with ThreadPoolExecutor(max_workers=MAX_WORKERS) as ex: |
| 644 | futures = [ex.submit(build_vector, d) for d in docs] |
| 645 | for f in tqdm( |