""" Stage 4 — embed (BGE-M3) + store (Qdrant), and Stage 5 — hybrid retrieval + rerank. WHY BGE-M3: it emits a DENSE vector (semantic meaning: "variance of a sum" ~ "Var(X+Y)") AND a SPARSE lexical vector (exact token/symbol overlap: matches "$\\sigma^2$", "binomial") from a single model. Dense alone misses exact symbol matches; sparse alone misses paraphrase. Fusing both (hybrid) is why this beats plain cosine on technical queries. WHY parent/child here: we embed small retrieval units (children, or childless parents) for precise matching, but store each unit's PARENT text in the payload so generation gets the whole problem+solution, not a fragment. """ from __future__ import annotations import os # torch and FlagEmbedding both ship an OpenMP runtime (libiomp5md.dll); on Windows the # duplicate triggers a native crash (0xC0000005). This makes the load robust. os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1") # IMPORT ORDER MATTERS: torch + FlagEmbedding native libs must load BEFORE qdrant_client, # or the second one to load segfaults (0xC0000005) on Windows. Proven by isolation test. import torch # noqa: F401 import FlagEmbedding # noqa: F401 import json from pathlib import Path from qdrant_client import QdrantClient, models DENSE_DIM = 1024 COLLECTION = "probability" _model = None _reranker = None def get_model(): """Lazy-load BGE-M3. fp16 only when a GPU is visible; on CPU use fp32. The web server hides the GPU (CUDA_VISIBLE_DEVICES="") so the full 8GB is free for Ollama's 7B generation — retrieval on CPU is light (~1-3s).""" global _model if _model is None: cuda = torch.cuda.device_count() > 0 # pass an explicit single device on CPU — auto-detect divides by GPU count (0 -> crash) _model = FlagEmbedding.BGEM3FlagModel( "BAAI/bge-m3", use_fp16=cuda, devices=None if cuda else "cpu") return _model def get_reranker(): """Lazy-load the bge-reranker-v2-m3 cross-encoder (fp16 on GPU, fp32 on CPU).""" global _reranker if _reranker is None: cuda = torch.cuda.device_count() > 0 _reranker = FlagEmbedding.FlagReranker( "BAAI/bge-reranker-v2-m3", use_fp16=cuda, devices=None if cuda else "cpu") return _reranker def embed(texts: list[str]) -> tuple[list[list[float]], list[dict]]: out = get_model().encode( texts, return_dense=True, return_sparse=True, return_colbert_vecs=False, batch_size=8, max_length=8192, ) dense = [v.tolist() for v in out["dense_vecs"]] sparse = [{int(k): float(v) for k, v in lw.items()} for lw in out["lexical_weights"]] return dense, sparse def _sparse_vec(d: dict[int, float]) -> models.SparseVector: return models.SparseVector(indices=list(d.keys()), values=list(d.values())) def contextualize(chapter: str, section: str, text: str) -> str: """Prepend a one-line location header ("Chapter 10: Generating Functions > Examples") to a chunk before it is embedded and reranked. WHY: worked examples state results in notation, not words -- e.g. Example 10.4 computes the Poisson mgf as "g(t) = ..." but never writes the phrase "moment generating function", so a query using that phrase out-ranks it onto longer prose-theory chunks. The chapter title ("Generating Functions") supplies the missing vocabulary, letting the example surface for the query that it actually answers. This is lightweight contextual retrieval -- no LLM, computed once at index time, and it also makes the generation context self-labelling.""" loc = " > ".join(p for p in (chapter, section) if p) return f"[{loc}]\n{text}" if loc else text def retrieval_units(chunks: list[dict]) -> list[dict]: """Units to embed: children, plus parents that have no children. Each unit is annotated with the parent text it should return. Front/back matter (title page, Contents, Preface, Index, appendix data tables) carries an empty chapter label and is dropped here. WHY: the Index is one 7600-token chunk that lists every term in the book, so it out-ranks real content on keyword queries (it surfaced as the top hit for "Bayes' theorem", burying the actual Bayes' Formula section and causing a refusal). None of this matter is explanatory content, so excluding it strictly improves retrieval precision.""" by_id = {c["id"]: c for c in chunks} parents_with_kids = {c["parent_id"] for c in chunks if c.get("parent_id")} units = [] for c in chunks: if not c["chapter"]: continue # front/back matter, not course content if c["is_parent"] and c["id"] in parents_with_kids: continue # represented by its children parent = by_id.get(c.get("parent_id")) if c.get("parent_id") else c u = dict(c) u["parent_text"] = parent["text"] u["parent_id"] = parent["id"] units.append(u) return units def build_collection(client: QdrantClient, name: str = COLLECTION) -> None: # Delete unconditionally rather than guarding on collection_exists(): points use # sequential int ids (0..N-1), so if a rebuild has FEWER units than before and the old # collection survives, upsert overwrites 0..N-1 and leaves the tail ids as stale points # (this once left the dropped book Index lingering and answerable). delete_collection is # a no-op tolerant call here. client.delete_collection(name) client.create_collection( name, vectors_config={"dense": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE)}, sparse_vectors_config={"sparse": models.SparseVectorParams(modifier=models.Modifier.IDF)}, ) def upsert_units(client: QdrantClient, units: list[dict], name: str = COLLECTION, pages: dict[str, int] | None = None) -> int: # Embed the contextualized child text (location header + chunk). The reranker scores # against the contextualized PARENT text (stored below), so both retrieval stages see # the chapter/section vocabulary the worked examples lack. pages = pages or {} embed_texts = [contextualize(u["chapter"], u["section"], u["text"]) for u in units] dense, sparse = embed(embed_texts) points = [] for i, u in enumerate(units): parent_text = contextualize(u["chapter"], u["section"], u["parent_text"]) points.append( models.PointStruct( id=i, vector={"dense": dense[i], "sparse": _sparse_vec(sparse[i])}, payload={ "chunk_id": u["id"], "parent_id": u["parent_id"], "type": u["type"], "chapter": u["chapter"], "section": u["section"], "source": u["source"], "text": u["text"], "parent_text": parent_text, "page": pages.get(u["parent_id"]), # 1-indexed PDF page for the UI viewer }, ) ) client.upsert(name, points) return len(points) def _dedupe_parents(hits: list[dict]) -> list[dict]: """Keep the best-scoring unit per parent so the LLM sees distinct problems.""" seen: dict[str, dict] = {} for h in hits: pid = h["parent_id"] if pid not in seen: seen[pid] = h return list(seen.values()) def rerank(query: str, hits: list[dict], *, top_k: int) -> list[dict]: """Cross-encoder rerank: score (query, parent_text) pairs, return top_k.""" if not hits: return [] # batch_size=8: FlagEmbedding's compute_score runs a throwaway "adjust batch size" # forward pass over the FIRST batch before the real loop — with the default (128) that # first batch is the whole pool, so every rerank did 2x the GPU work. A small batch # caps the waste at 8 pairs. Scores are identical (measured: 51 pairs 2.0s -> 1.06s). scores = get_reranker().compute_score( [[query, h["parent_text"]] for h in hits], normalize=True, batch_size=8) if not isinstance(scores, list): scores = [scores] for h, s in zip(hits, scores): h["rerank_score"] = float(s) hits.sort(key=lambda h: h["rerank_score"], reverse=True) return hits[:top_k] def _hybrid_candidates(client: QdrantClient, query: str, cand: int, name: str) -> list[dict]: """One dense+sparse RRF retrieval → raw candidate payloads (pre-dedupe, pre-rerank).""" dense, sparse = embed([query]) res = client.query_points( name, prefetch=[ models.Prefetch(query=dense[0], using="dense", limit=cand), models.Prefetch(query=_sparse_vec(sparse[0]), using="sparse", limit=cand), ], query=models.FusionQuery(fusion=models.Fusion.RRF), limit=cand, with_payload=True, ) return [{"score": p.score, **p.payload} for p in res.points] def search(client: QdrantClient, query: str, *, top_k: int = 5, mode: str = "rerank", name: str = COLLECTION) -> list[dict]: """mode: 'dense' (vector-only), 'hybrid' (dense+sparse RRF), or 'rerank' (hybrid candidates re-scored by the cross-encoder).""" # Candidate pool feeds the cross-encoder, whose cost is linear (~40ms/pair on the # 4060, shared with Ollama) — the pool size IS the time-to-first-token. 80 candidates # cost ~5s/query; A/B on the 45-question golden set showed 24 is metric-identical # (Hit@1/3/5, MRR 0.881, same misses), so the old "big pool recovers vocabulary # mismatch" worry didn't survive measurement. cand = max(top_k * 4, 24) if mode == "dense": dense, _ = embed([query]) res = client.query_points(name, query=dense[0], using="dense", limit=cand, with_payload=True) hits = _dedupe_parents([{"score": p.score, **p.payload} for p in res.points]) else: hits = _dedupe_parents(_hybrid_candidates(client, query, cand, name)) if mode == "rerank": return rerank(query, hits, top_k=top_k) return hits[:top_k] def search_multi(client: QdrantClient, queries: list[str], *, top_k: int = 5, cand: int = 40, name: str = COLLECTION) -> list[dict]: """Multi-query retrieval for comparison / multi-part questions. WHY: one query embedding can't sit near two distributions at once, so a comparison retrieves only one side. Each sub-query is retrieved AND reranked against ITSELF (a chunk about only the geometric distribution must be scored on "geometric variance", not on "compare binomial and geometric" — the latter buries single-topic chunks, which is exactly what killed the union-then-rerank approach). The per-query rankings are then round-robin merged, so every sub-query is guaranteed representation in the final top_k and both sides surface. LLM-free by design — the caller (generate.decompose) makes the sub-queries.""" lanes = [] for q in queries: hits = _dedupe_parents(_hybrid_candidates(client, q, cand, name)) lanes.append(rerank(q, hits, top_k=top_k)) # scored against its OWN query merged, seen = [], set() for i in range(top_k): for lane in lanes: if i < len(lane) and lane[i]["parent_id"] not in seen: seen.add(lane[i]["parent_id"]) merged.append(lane[i]) if len(merged) >= top_k: return merged return merged def load_chunks(path: Path) -> list[dict]: return json.loads(Path(path).read_text(encoding="utf-8"))