Add semantic search

.dockerignore

@@ -60,4 +60,6 @@ src/backend/search_engine/tests/
 src/backend/search_engine/index_builder/build/
 src/backend/search_engine/index_builder/data/
 src/backend/search_engine/index/bin/
-src/backend/search_engine/models/neuspell-scrnn-probwordnoise/
+src/backend/search_engine/models/nomic-embed-text/
+src/backend/search_engine/models/neuspell-scrnn-probwordnoise/
+src/backend/search_engine/models/IVFPQ.faiss

.gitignore

@@ -233,10 +233,14 @@ dist-ssr
 
 *.tsv
 *.gz
+/src/backend/search_engine/models/nomic-embed-text/
 # allow test data files as they are small and necessary for tests
 !src/backend/search_engine/index_builder/test_data/*.tsv
 !src/backend/search_engine/index_builder/test_data/*.gz
 
 /src/backend/search_engine/index_builder/data/
 !src/backend/search_engine/index_builder/data/data.md
 /src/backend/search_engine/index/bin/
+src/backend/search_engine/models/IVFPQ.faiss
+src/backend/search_engine/models/neuspell-scrnn-probwordnoise
+memory_log.txt

README.md

@@ -21,7 +21,6 @@ Seekr consists of several core subsystems working together:
 * Node.js
 * npm
 * Docker (for containerized integration/unit tests & deployment)
-* LFS (downloading ML models from GitHub)
 * CMake (building and compiling the CPP components)
 * Just (command runner)
 
@@ -59,15 +58,15 @@ cd src && uv run --project backend python -m backend.search_engine.scripts.downl
 
 Build the index with a memory limit:
 ```bash
-just build-index <memory-limit>
+just build-index <memory-limit> [<max-docs>]
 ```
 
 Start both frontend and backend in dev environment (`http://localhost:8080`):
 ```bash
-just local <uvicorn-args>
+just local [<uvicorn-args>]
 ```
 
 Query Search Engine from CLI:
 ```bash
 cd src && uv run --project backend python -m backend.search_engine.scripts.query --query <query> --limit <limit>
-```
+```

src/backend/api/v1/app.py

@@ -1,25 +1,37 @@
-from typing import Annotated
 import os
 from contextlib import asynccontextmanager
-from fastapi import FastAPI, Query, HTTPException, status
-from fastapi.middleware.cors import CORSMiddleware
+from typing import Annotated
 
-from backend.logging_config import setup_logging, get_logger
-from backend.search_engine.models.index import SearchResults
+from backend.logging_config import get_logger, setup_logging
+from backend.memory_tracer import LOG_FILE
+from backend.search_engine.error_handling import InvalidOperatorError
 from backend.search_engine.index.index_loader import get_index
+from backend.search_engine.models.index import SearchResults
 from backend.search_engine.query.query_engine import QueryEngine
-from backend.search_engine.error_handling import InvalidOperatorError
+from backend.search_engine.semantic_search.embedding_model import (
+    MAX_QUERY_LENGTH,
+    get_embedding_model,
+)
+from backend.search_engine.semantic_search.train_vector_index import train_or_load_ivfpq
 from backend.search_engine.spell_correction.spell_corrector import get_spell_corrector
+from fastapi import FastAPI, HTTPException, Query, status
+from fastapi.middleware.cors import CORSMiddleware
 
 setup_logging(level=os.getenv("LOG_LEVEL", "INFO"))
 logger = get_logger(__name__)
 
+if LOG_FILE.exists():
+    LOG_FILE.unlink()
+
 
+# preload all large models at the beginning
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     logger.debug("Starting up...")
     app.state.inverted_index = get_index()
     app.state.spell_corrector = get_spell_corrector()
+    app.state.embedding_model = get_embedding_model()
+    app.state.vector_index = train_or_load_ivfpq()
     yield
     logger.debug("Shutting down...")
 
@@ -37,7 +49,10 @@ async def lifespan(app: FastAPI):
 
 @app.get("/search", response_model=SearchResults)
 async def search(
-    q: Annotated[str, Query(min_length=1, max_length=50, description="Search query")],
+    q: Annotated[
+        str,
+        Query(min_length=1, max_length=MAX_QUERY_LENGTH, description="Search query"),
+    ],
     limit: Annotated[
         int, Query(ge=1, le=500, description="Maximum number of results")
     ] = 10,

src/backend/bindings/utils.cpp

@@ -147,8 +147,10 @@ std::vector<std::string> normalize_search_query(const std::string& text) {
 
 struct Metadata {
     uint32_t num_docs = 0;
-    double avg_doc_length = 0.0;
-    std::unordered_map<uint32_t, uint32_t> doc_lengths;
+    double avg_title_length = 0.0;
+    double avg_body_length = 0.0;
+    std::unordered_map<uint32_t, uint32_t> title_lengths;
+    std::unordered_map<uint32_t, uint32_t> body_lengths;
 
     void load(const std::string& path) {
         std::cout << "Metadata: " << path << std::endl;
@@ -157,21 +159,33 @@ struct Metadata {
         if (in.peek() == EOF) throw std::runtime_error("Metadata file is empty");
 
         in.read(reinterpret_cast<char*>(&num_docs), sizeof(num_docs));
-        in.read(reinterpret_cast<char*>(&avg_doc_length), sizeof(avg_doc_length));
+        in.read(reinterpret_cast<char*>(&avg_title_length), sizeof(avg_title_length));
+        in.read(reinterpret_cast<char*>(&avg_body_length), sizeof(avg_body_length));
 
         while (in.peek() != EOF) {
-            uint32_t doc_id, length;
+            uint32_t doc_id, t_len, b_len;
             if (!in.read(reinterpret_cast<char*>(&doc_id), sizeof(doc_id))) break;
-            if (!in.read(reinterpret_cast<char*>(&length), sizeof(length))) break;
-            doc_lengths[doc_id] = length;
+            if (!in.read(reinterpret_cast<char*>(&t_len), sizeof(t_len))) break;
+            if (!in.read(reinterpret_cast<char*>(&b_len), sizeof(b_len))) break;
+            title_lengths[doc_id] = t_len;
+            body_lengths[doc_id] = b_len;
         }
     }
 
-    uint32_t get_doc_length(uint32_t doc_id) const {
-        auto it = doc_lengths.find(doc_id);
-        if (it == doc_lengths.end()) return 0;
+    uint32_t get_title_length(uint32_t doc_id) const {
+        auto it = title_lengths.find(doc_id);
+        if (it == title_lengths.end()) return 0;
         return it->second;
     }
+
+    uint32_t get_body_length(uint32_t doc_id) const {
+        auto it = body_lengths.find(doc_id);
+        if (it == body_lengths.end()) return 0;
+        return it->second;
+    }
+
+    // fallback for compatibility
+    uint32_t get_doc_length(uint32_t doc_id) const { return get_body_length(doc_id); }
 };
 
 struct PostingList {
@@ -985,8 +999,12 @@ PYBIND11_MODULE(_core, m) {
 
     py::class_<Metadata>(m, "Metadata")
         .def_readonly("num_docs", &Metadata::num_docs)
-        .def_readonly("avg_doc_length", &Metadata::avg_doc_length)
-        .def_readonly("doc_lengths", &Metadata::doc_lengths)
+        .def_readonly("avg_title_length", &Metadata::avg_title_length)
+        .def_readonly("avg_body_length", &Metadata::avg_body_length)
+        .def_readonly("title_lengths", &Metadata::title_lengths)
+        .def_readonly("body_lengths", &Metadata::body_lengths)
+        .def("get_title_length", &Metadata::get_title_length, py::arg("doc_id"))
+        .def("get_body_length", &Metadata::get_body_length, py::arg("doc_id"))
         .def("get_doc_length", &Metadata::get_doc_length, py::arg("doc_id"));
 
     py::class_<DocStore>(m, "DocStore")

src/backend/memory_tracer.py

@@ -0,0 +1,89 @@
+import gc
+import os
+import time
+from functools import wraps
+from pathlib import Path
+
+import numpy as np
+import psutil  # type: ignore [import-untyped]
+
+LOG_FILE = Path(__file__).resolve().parent / "memory_log.txt"
+
+MB = 1024**2
+
+
+def _object_size_mb(obj) -> tuple[str, float]:
+    """Measure the actual memory footprint of a returned object."""
+    import torch
+
+    # torch model (nn.Module or dataclass containing one)
+    model = None
+    if isinstance(obj, torch.nn.Module):
+        model = obj
+    elif hasattr(obj, "model") and isinstance(getattr(obj, "model"), torch.nn.Module):
+        model = obj.model
+
+    if model is not None:
+        params = sum(p.numel() * p.element_size() for p in model.parameters())
+        buffers = sum(b.numel() * b.element_size() for b in model.buffers())
+        total = (params + buffers) / MB
+        return "torch_model", total
+
+    # faiss index
+    try:
+        import faiss  # type: ignore [import-untyped]
+
+        if isinstance(obj, faiss.Index):
+            # write to a temporary buffer to measure serialized size
+            writer = faiss.VectorIOWriter()
+            faiss.write_index(obj, writer)
+            total = len(faiss.vector_to_array(writer.data)) / MB
+            return "faiss_index", total
+    except ImportError:
+        pass
+
+    # numpy array
+    if isinstance(obj, np.ndarray):
+        return "numpy", obj.nbytes / MB
+
+    return "unknown", 0.0
+
+
+def trace_memory(func):
+    """Universal memory tracer: measures RSS delta + actual object size."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        proc = psutil.Process(os.getpid())
+        gc.collect()
+
+        mem_before = proc.memory_info().rss
+        start_time = time.perf_counter()
+
+        result = func(*args, **kwargs)
+
+        duration = time.perf_counter() - start_time
+        gc.collect()
+        if not LOG_FILE.exists() or LOG_FILE.stat().st_size == 0:
+            with open(LOG_FILE, "a") as f:
+                f.write(f"Process PID: {os.getpid()}\n")
+                f.write("------------------------------------\n\n")
+
+        mem_after = proc.memory_info().rss
+
+        obj_type, obj_size = _object_size_mb(result)
+
+        with open(LOG_FILE, "a") as f:
+            f.write(f"--- {func.__name__} ---\n")
+            f.write(f"Time:          {duration:.4f}s\n")
+            f.write(f"RSS before:    {mem_before / MB:.2f} MB\n")
+            f.write(f"RSS after:     {mem_after / MB:.2f} MB\n")
+            f.write(f"RSS delta:     {(mem_after - mem_before) / MB:.2f} MB\n")
+            if obj_size > 0:
+                f.write(f"Object type:   {obj_type}\n")
+                f.write(f"Object size:   {obj_size:.2f} MB\n")
+            f.write("------------------------------------\n\n")
+
+        return result
+
+    return wrapper

src/backend/mypy.ini

@@ -1,2 +1,2 @@
 [mypy]
-exclude = scripts/
+exclude = scripts/|search_engine/models/nomic-embed-text/

src/backend/pyproject.toml

@@ -10,6 +10,11 @@ dependencies = [
     "requests>=2.32.5",
     "tqdm>=4.67.1",
     "neuspell>=1.0.0",
+    "polars>=1.38.1",
+    "einops>=0.8.2", # for huggingface embedding model
+    "faiss-cpu>=1.13.2",
+    "sentence-transformers>=5.2.3",
+    "psutil>=7.2.2",
 ]
 
 [dependency-groups]

src/backend/ruff.toml

@@ -1 +1 @@
-exclude = ["*.ipynb"]
+exclude = ["*.ipynb", "search_engine/models/nomic-embed-text"]

src/backend/search_engine/index/index_loader.py

@@ -1,8 +1,11 @@
 import os
 import time
 from functools import lru_cache
+
 from cpp_utils import InvertedIndex  # type: ignore [import-untyped]
+
 from backend.logging_config import get_logger
+from backend.memory_tracer import trace_memory
 
 logger = get_logger(__name__)
 
@@ -17,6 +20,7 @@ def _index_path() -> str:
 
 
 @lru_cache(maxsize=1)
+@trace_memory
 def get_index() -> InvertedIndex:
     logger.debug(f"Loading inverted index from {_index_path()}...")
     start = time.perf_counter()