From b32e466140bb8db442c4289fedfb7bd7f9fda982 Mon Sep 17 00:00:00 2001 From: Matt Weiden <538456+mweiden@users.noreply.github.com> Date: Tue, 12 Aug 2025 22:23:16 -0700 Subject: [PATCH] search: rename InvertedIndex to SearchIndex --- src/app.py | 23 +-- src/pickle_store.py | 10 +- .../{inverted_index.py => search_index.py} | 68 ++++++-- src/search/tests/test_inverted_index.py | 101 ------------ src/search/tests/test_search_index.py | 146 ++++++++++++++++++ src/tests/test_pickle_store_mps.py | 4 +- src/web_crawler/tests/test_web_scraper.py | 4 +- src/web_crawler_cron.py | 14 +- 8 files changed, 232 insertions(+), 138 deletions(-) rename src/search/{inverted_index.py => search_index.py} (57%) delete mode 100644 src/search/tests/test_inverted_index.py create mode 100644 src/search/tests/test_search_index.py diff --git a/src/app.py b/src/app.py index 7da7fd0..62a8c88 100644 --- a/src/app.py +++ b/src/app.py @@ -6,7 +6,8 @@ from autocomplete.mermaid import Mermaid from autocomplete.subgraph_cache_trie import SubgraphCacheTrie -from search.inverted_index import InvertedIndex +from search.search_index import SearchIndex +from sentence_transformers import CrossEncoder from pickle_store import PickleStore from util import add_file_handler, get_static_file from env import ( @@ -31,14 +32,15 @@ HTML_TRIE = get_static_file("trie.html") # inverted index config -INVERTED_INDEX_STORAGE = PickleStore(INVERTED_INDEX_STORAGE_PATH) -INVERTED_INDEX = INVERTED_INDEX_STORAGE.get_latest(InvertedIndex).artifact +SEARCH_INDEX_STORAGE = PickleStore(INVERTED_INDEX_STORAGE_PATH) +SEARCH_INDEX = SEARCH_INDEX_STORAGE.get_latest(SearchIndex).artifact +SEARCH_INDEX.set_cross_encoder(CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")) # autocomplete config TRIE_STORAGE = PickleStore(TRIE_STORAGE_PATH) try: AUTOCOMPLETE_TRIE = TRIE_STORAGE.get_latest(SubgraphCacheTrie).artifact -except AttributeError as e: +except AttributeError: AUTOCOMPLETE_TRIE = SubgraphCacheTrie() MERMAID = Mermaid() @@ -63,16 +65,17 @@ def search(): query = request.form["query"] analytics_logger.info(query) app.logger.info(f"Received query: {query}") - search_results = INVERTED_INDEX.top_k(query) + search_results = SEARCH_INDEX.top_k(query) return [asdict(result) for result in search_results] @app.route("/inverted-index/load", methods=["POST"]) -def load_inverted_index(): - global INVERTED_INDEX - inverted_index_blob = INVERTED_INDEX_STORAGE.get_latest(InvertedIndex) - INVERTED_INDEX = inverted_index_blob.artifact - app.logger.info(f"Loaded new inverted index: {inverted_index_blob.file_path}") +def load_search_index(): + global SEARCH_INDEX + search_index_blob = SEARCH_INDEX_STORAGE.get_latest(SearchIndex) + SEARCH_INDEX = search_index_blob.artifact + SEARCH_INDEX.set_cross_encoder(CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")) + app.logger.info(f"Loaded new search index: {search_index_blob.file_path}") return dict(status="OK") diff --git a/src/pickle_store.py b/src/pickle_store.py index f2aa6fd..e0018af 100644 --- a/src/pickle_store.py +++ b/src/pickle_store.py @@ -6,13 +6,13 @@ from datetime import datetime from autocomplete.subgraph_cache_trie import SubgraphCacheTrie -from search.inverted_index import InvertedIndex +from search.search_index import SearchIndex @dataclass class Artifact: file_path: str - artifact: SubgraphCacheTrie | InvertedIndex + artifact: SubgraphCacheTrie | SearchIndex class PickleStore: @@ -22,7 +22,7 @@ def __init__(self, dir: str): raise FileNotFoundError(f"Directory does not exist: {dir}") self.dir = dir - def save(self, obj: SubgraphCacheTrie | InvertedIndex) -> str: + def save(self, obj: SubgraphCacheTrie | SearchIndex) -> str: now = datetime.now() formatted_date = now.strftime("%Y%m%d%H%M%S") file_prefix = self._file_prefix(type(obj)) @@ -31,7 +31,7 @@ def save(self, obj: SubgraphCacheTrie | InvertedIndex) -> str: pickle.dump(obj, file) return pkl_file_name - def get_latest(self, typ: SubgraphCacheTrie | InvertedIndex) -> Optional[Artifact]: + def get_latest(self, typ: SubgraphCacheTrie | SearchIndex) -> Optional[Artifact]: file_prefix = self._file_prefix(typ) matching_files = glob.glob(f"{self.dir}/{file_prefix}_*.pkl") if not matching_files: @@ -81,5 +81,5 @@ def _cpu_deserialize(storage, location): # type: ignore def _file_prefix(self, typ: Type) -> str: if typ == SubgraphCacheTrie: return "trie" - elif typ == InvertedIndex: + elif typ == SearchIndex: return "inverted_index" diff --git a/src/search/inverted_index.py b/src/search/search_index.py similarity index 57% rename from src/search/inverted_index.py rename to src/search/search_index.py index ee7401e..6648024 100644 --- a/src/search/inverted_index.py +++ b/src/search/search_index.py @@ -1,9 +1,9 @@ from dataclasses import dataclass, field -from collections import defaultdict +from collections import defaultdict, OrderedDict import hnswlib import numpy as np -from sentence_transformers import SentenceTransformer +from sentence_transformers import CrossEncoder, SentenceTransformer from web_crawler.node import Node from search.tokenizer import tokenize @@ -17,17 +17,21 @@ class SearchResult: score: float | None = field(default=None) -class InvertedIndex: +class SearchIndex: def __init__( self, model: SentenceTransformer | None = None, model_name: str = "sentence-transformers/paraphrase-MiniLM-L3-v2", + cross_encoder: CrossEncoder | None = None, + cross_encoder_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2", + cache_size: int = 1_000, ): self._inverted_index: dict[str, list[tuple[str, int]]] = defaultdict(list) self._words_per_doc: dict[str, int] = defaultdict(int) self._doc_id_to_url: dict[str, str] = dict() self._doc_id_to_title: dict[str, str] = dict() + self._doc_id_to_text: dict[str, str] = dict() self._model = model if model is not None else SentenceTransformer(model_name) embedding_dim = self._model.get_sentence_embedding_dimension() self._hnsw = hnswlib.Index(space="cosine", dim=embedding_dim) @@ -35,6 +39,14 @@ def __init__( self._id_to_label: dict[str, int] = {} self._label_to_id: dict[int, str] = {} self._next_label = 0 + if cross_encoder is not None: + self._cross_encoder = cross_encoder + elif cross_encoder_name is not None: + self._cross_encoder = CrossEncoder(cross_encoder_name) + else: + self._cross_encoder = None + self._score_cache: OrderedDict[tuple[str, str], float] = OrderedDict() + self._cache_size = cache_size @property def total_docs(self): @@ -48,6 +60,7 @@ def insert(self, doc: Node) -> None: self._doc_id_to_url[doc.id] = doc.url self._doc_id_to_title[doc.id] = doc.title if doc.text is not None: + self._doc_id_to_text[doc.id] = doc.text embedding = self._model.encode(doc.text).astype(np.float32) if self._next_label >= self._hnsw.get_max_elements(): self._hnsw.resize_index(self._next_label * 2) @@ -83,23 +96,56 @@ def search(self, word: str) -> list[SearchResult]: def num_words_in_doc(self, doc_id: str) -> int: return self._words_per_doc[doc_id] - def top_k(self, query: str, k: int = 10) -> list[SearchResult]: + def set_cross_encoder(self, cross_encoder: CrossEncoder) -> None: + self._cross_encoder = cross_encoder + + def top_k( + self, query: str, k: int = 10, candidates: int | None = None + ) -> list[SearchResult]: if self._next_label == 0: return [] - k = min(k, self._next_label) + candidates = self._next_label if candidates is None else candidates + candidates = min(max(k, candidates), self._next_label) query_embedding = self._model.encode(query).astype(np.float32) - self._hnsw.set_ef(max(k, 10)) - labels, distances = self._hnsw.knn_query(query_embedding, k) + self._hnsw.set_ef(max(candidates, 10)) + labels, distances = self._hnsw.knn_query(query_embedding, candidates) results = [] - for label, distance in zip(labels[0], distances[0]): + for label in labels[0]: doc_id = self._label_to_id[label] - score = 1.0 - float(distance) results.append( SearchResult( id=doc_id, url=self._doc_id_to_url[doc_id], title=self._doc_id_to_title[doc_id], - score=score, + score=None, ) ) - return results + if self._cross_encoder is None: + for result, distance in zip(results, distances[0]): + result.score = 1.0 - float(distance) + return results[:k] + missing: list[SearchResult] = [] + pairs: list[tuple[str, str]] = [] + for result in results: + text = self._doc_id_to_text.get(result.id, "") + key = (query, result.id) + if key in self._score_cache: + self._score_cache.move_to_end(key) + result.score = self._score_cache[key] + else: + missing.append(result) + pairs.append((query, text)) + if pairs: + scores = self._cross_encoder.predict(pairs) + for result, score in zip(missing, scores): + result.score = float(score) + key = (query, result.id) + self._score_cache[key] = result.score + self._score_cache.move_to_end(key) + if len(self._score_cache) > self._cache_size: + self._score_cache.popitem(last=False) + results.sort( + key=lambda r: r.score if r.score is not None else float("-inf"), + reverse=True, + ) + return results[:k] diff --git a/src/search/tests/test_inverted_index.py b/src/search/tests/test_inverted_index.py deleted file mode 100644 index 5d943ac..0000000 --- a/src/search/tests/test_inverted_index.py +++ /dev/null @@ -1,101 +0,0 @@ -import numpy as np -import pytest -import requests -from sentence_transformers import SentenceTransformer - -from web_crawler.node import Node -from search.inverted_index import InvertedIndex, SearchResult -from search.tokenizer import tokenize - - -@pytest.fixture -def node0() -> Node: - return Node( - raw_url="https://some.page.on.the.internet", - text=( - "In publishing and graphic design, Lorem ipsum is a placeholder text " - "commonly used to demonstrate the visual form of a document or a typeface " - "without relying on meaningful content iterators." - ), - title="lorem ipsum", - ) - - -@pytest.fixture -def node1() -> Node: - return Node( - raw_url="https://some.other.page.on.the.internet", - text=( - "For clarity and correctness in type hints, you should use Iterable and/or " - "Iterator from the typing module to annotate functions that return iterators." - ), - title="iterable/iterator", - ) - - -@pytest.fixture -def inverted_index(node0, node1) -> InvertedIndex: - class DummyModel: - KEYWORDS0 = {"lorem", "ipsum"} - KEYWORDS1 = {"type", "hint", "iterators", "iterator"} - - def encode(self, text: str) -> np.ndarray: - tokens = text.lower().split() - vec = np.array( - [ - sum(t in self.KEYWORDS0 for t in tokens), - sum(t in self.KEYWORDS1 for t in tokens), - ], - dtype=np.float32, - ) - if not vec.any(): - vec = np.array([len(tokens), 0.0], dtype=np.float32) - return vec - - def get_sentence_embedding_dimension(self) -> int: - return 2 - - model = DummyModel() - inverted_index = InvertedIndex(model=model) - inverted_index.insert(node0) - inverted_index.insert(node1) - return inverted_index - - -@pytest.fixture -def search_result0(node0) -> SearchResult: - return SearchResult(node0.id, node0.url, node0.title) - - -@pytest.fixture -def search_result1(node1) -> SearchResult: - return SearchResult(node1.id, node1.url, node1.title) - - -def test_insert(inverted_index, search_result0, search_result1): - assert list(tokenize("ipsum")) == ["ipsum"] - assert list(tokenize("iterators")) == ["iterators"] - assert inverted_index.search("ipsum") == [search_result0] - assert inverted_index.search("iterators") == [search_result0, search_result1] - - -def test_num_words_in_doc(inverted_index, node0, node1): - node0_tokens = list(tokenize(node0.text)) - node1_tokens = list(tokenize(node1.text)) - assert inverted_index.num_words_in_doc(node0.id) == len(node0_tokens) - assert inverted_index.num_words_in_doc(node1.id) == len(node1_tokens) - - -def test_top_k_cosine(inverted_index, node0, node1): - results = inverted_index.top_k("lorem ipsum") - assert results[0].id == node0.id - - results = inverted_index.top_k("type hint iterators") - assert [r.id for r in results[:2]] == [node1.id, node0.id] - - -def test_paraphrase_minilm_model_simple_example(): - pytest.skip( - "requires downloading a large SentenceTransformer model", - allow_module_level=False, - ) diff --git a/src/search/tests/test_search_index.py b/src/search/tests/test_search_index.py new file mode 100644 index 0000000..1f85262 --- /dev/null +++ b/src/search/tests/test_search_index.py @@ -0,0 +1,146 @@ +import numpy as np +import pytest + +from web_crawler.node import Node +from search.search_index import SearchIndex, SearchResult +from search.tokenizer import tokenize + + +@pytest.fixture +def node0() -> Node: + return Node( + raw_url="https://some.page.on.the.internet", + text=( + "In publishing and graphic design, Lorem ipsum is a placeholder text " + "commonly used to demonstrate the visual form of a document or a typeface " + "without relying on meaningful content iterators." + ), + title="lorem ipsum", + ) + + +@pytest.fixture +def node1() -> Node: + return Node( + raw_url="https://some.other.page.on.the.internet", + text=( + "For clarity and correctness in type hints, you should use Iterable and/or " + "Iterator from the typing module to annotate functions that return iterators." + ), + title="iterable/iterator", + ) + + +@pytest.fixture +def search_index(node0, node1) -> SearchIndex: + class DummyModel: + KEYWORDS0 = {"lorem", "ipsum"} + KEYWORDS1 = {"type", "hint", "iterators", "iterator"} + + def encode(self, text: str) -> np.ndarray: + tokens = text.lower().split() + vec = np.array( + [ + sum(t in self.KEYWORDS0 for t in tokens), + sum(t in self.KEYWORDS1 for t in tokens), + ], + dtype=np.float32, + ) + if not vec.any(): + vec = np.array([len(tokens), 0.0], dtype=np.float32) + return vec + + def get_sentence_embedding_dimension(self) -> int: + return 2 + + model = DummyModel() + index = SearchIndex(model=model, cross_encoder=None, cross_encoder_name=None) + index.insert(node0) + index.insert(node1) + return index + + +@pytest.fixture +def search_result0(node0) -> SearchResult: + return SearchResult(node0.id, node0.url, node0.title) + + +@pytest.fixture +def search_result1(node1) -> SearchResult: + return SearchResult(node1.id, node1.url, node1.title) + + +def test_insert(search_index, search_result0, search_result1): + assert list(tokenize("ipsum")) == ["ipsum"] + assert list(tokenize("iterators")) == ["iterators"] + assert search_index.search("ipsum") == [search_result0] + assert search_index.search("iterators") == [search_result0, search_result1] + + +def test_num_words_in_doc(search_index, node0, node1): + node0_tokens = list(tokenize(node0.text)) + node1_tokens = list(tokenize(node1.text)) + assert search_index.num_words_in_doc(node0.id) == len(node0_tokens) + assert search_index.num_words_in_doc(node1.id) == len(node1_tokens) + + +def test_top_k_cosine(search_index, node0, node1): + results = search_index.top_k("lorem ipsum") + assert results[0].id == node0.id + + results = search_index.top_k("type hint iterators") + assert [r.id for r in results[:2]] == [node1.id, node0.id] + + +def test_cross_encoder_rerank_and_cache(node0, node1): + class DummyModel: + KEYWORDS0 = {"lorem", "ipsum"} + KEYWORDS1 = {"type", "hint", "iterators", "iterator"} + + def encode(self, text: str) -> np.ndarray: + tokens = text.lower().split() + vec = np.array( + [ + sum(t in self.KEYWORDS0 for t in tokens), + sum(t in self.KEYWORDS1 for t in tokens), + ], + dtype=np.float32, + ) + if not vec.any(): + vec = np.array([len(tokens), 0.0], dtype=np.float32) + return vec + + def get_sentence_embedding_dimension(self) -> int: + return 2 + + class DummyCrossEncoder: + def __init__(self): + self.call_count = 0 + + def predict(self, pairs): + self.call_count += len(pairs) + scores = [] + for _, text in pairs: + scores.append(2.0 if "lorem" in text.lower() else 1.0) + return scores + + model = DummyModel() + cross_encoder = DummyCrossEncoder() + index = SearchIndex(model=model, cross_encoder=cross_encoder) + index.insert(node0) + index.insert(node1) + + results = index.top_k("type hint iterators", k=2, candidates=2) + assert [r.id for r in results] == [node0.id, node1.id] + assert cross_encoder.call_count == 2 + + # second call uses cache, so cross_encoder isn't invoked again + index.top_k("type hint iterators", k=2, candidates=2) + assert cross_encoder.call_count == 2 + + +def test_paraphrase_minilm_model_simple_example(): + pytest.skip( + "requires downloading a large SentenceTransformer model", + allow_module_level=False, + ) diff --git a/src/tests/test_pickle_store_mps.py b/src/tests/test_pickle_store_mps.py index cbb9c60..24784e6 100644 --- a/src/tests/test_pickle_store_mps.py +++ b/src/tests/test_pickle_store_mps.py @@ -5,7 +5,7 @@ import torch from pickle_store import PickleStore -from search.inverted_index import InvertedIndex +from search.search_index import SearchIndex class Dummy: @@ -27,6 +27,6 @@ def test_get_latest_handles_mps_pickles(): file_path = os.path.join(tmpdir, "inverted_index_19700101000000.pkl") _create_mps_pickle(file_path, Dummy(42)) store = PickleStore(tmpdir) - artifact = store.get_latest(InvertedIndex) + artifact = store.get_latest(SearchIndex) assert artifact is not None assert artifact.artifact.tensor.item() == 42 diff --git a/src/web_crawler/tests/test_web_scraper.py b/src/web_crawler/tests/test_web_scraper.py index 0a78d66..c741897 100644 --- a/src/web_crawler/tests/test_web_scraper.py +++ b/src/web_crawler/tests/test_web_scraper.py @@ -1,9 +1,9 @@ import pytest -pytest.skip("requires Chrome driver and network access", allow_module_level=True) - from web_crawler.web_scraper import WebScraper +pytest.skip("requires Chrome driver and network access", allow_module_level=True) + scraper = WebScraper() scraper.navigate_sync("https://mweiden.github.io") diff --git a/src/web_crawler_cron.py b/src/web_crawler_cron.py index ad0d45e..4d8b1bc 100644 --- a/src/web_crawler_cron.py +++ b/src/web_crawler_cron.py @@ -8,7 +8,7 @@ from web_crawler.web_scraper import WebScraper from web_crawler.node import Node -from search.inverted_index import InvertedIndex +from search.search_index import SearchIndex from pickle_store import PickleStore from env import INVERTED_INDEX_STORAGE_PATH @@ -25,7 +25,7 @@ def seconds_since_program_start(): async def worker( worker_id: int, - inverted_index: InvertedIndex, + search_index: SearchIndex, max_depth: int, visited: set[str], netloc_last_visited_at: dict[str, int], @@ -47,7 +47,7 @@ async def worker( node.text = scraper.extract_rendered_text() node.title = scraper.driver.title - inverted_index.insert(node) + search_index.insert(node) links = scraper.find_all_links() @@ -67,7 +67,7 @@ async def worker( await queue.put(link_node) logger.info( - f"index_bytes={sys.getsizeof(inverted_index._inverted_index)} " + f"index_bytes={sys.getsizeof(search_index._inverted_index)} " f"queue_size={queue.qsize()} " f"depth={node.depth} " f"priority={node.priority} " @@ -94,7 +94,7 @@ async def main(): num_workers = 32 seed_url = "https://news.ycombinator.com" - inverted_index = InvertedIndex() + search_index = SearchIndex() pickle_store = PickleStore(f"../{INVERTED_INDEX_STORAGE_PATH}") seed_node = Node(seed_url) @@ -111,7 +111,7 @@ async def main(): asyncio.create_task( worker( worker_id=i, - inverted_index=inverted_index, + search_index=search_index, max_depth=max_depth, visited=visited, netloc_last_visited_at=netloc_last_visited_at, @@ -128,7 +128,7 @@ async def main(): await asyncio.gather(*workers, return_exceptions=True) - pickle_store.save(inverted_index) + pickle_store.save(search_index) if __name__ == "__main__":