| """Where blocking gets its embedding matrices from. |
| |
| Two interchangeable sources, selected by EMB["source"] (env M3_EMB_SOURCE): |
| |
| "internal" the shard cache written by embed.py / encoder.py (default) |
| "hybrid_bridge" the pre-computed multilingual-e5-large matrices already on |
| disk in /scratch/user4/dai/hybrid_bridge |
| |
| The hybrid_bridge matrices were produced with, verified on 2026-09-26 by |
| re-encoding 12 records and comparing against the stored vectors: |
| |
| text = "name: {business_name} | addr: {business_address} | country: {country}" |
| S1 = "query: " + text (mean cos vs stored 0.99902) |
| S2/S3 = "passage: " + text (mean cos vs stored 0.99782) |
| model = intfloat/multilingual-e5-large (MIT), 1024-d, fp16, L2-normalised |
| max_len= 96 (measured: p99 = 58 tokens, so nothing is truncated) |
| |
| Row order is the exact row order of the corresponding parquet in |
| dataset/dataset/hf/{split}_p/, which is itself the row order of the original |
| TSVs. load() asserts the row counts, so a silent misalignment cannot pass. |
| |
| MEASURED TRADE-OFF (20k real train queries, full 10.32M pool, country filter on, |
| dense channel only) -- see README "Embedding source": |
| |
| e5-large (hybrid_bridge) @20 = 0.9559 @30 = 0.9619 @50 = 0.9681 |
| e5-small (internal) @20 = 0.9652 |
| |
| The 1024-d e5-large matrices are 21 GB and score BELOW the 384-d e5-small ones |
| at the same k. Re-encoding the S1 side to make the prefixes symmetric was tested |
| and is worse still (passage:/passage: @20 = 0.9319), so the asymmetry is not the |
| cause. Use "hybrid_bridge" only with a larger emb_k, or to union the two dense |
| channels; it is not a drop-in improvement. |
| """ |
| from __future__ import annotations |
|
|
| import os |
| import numpy as np |
|
|
| import config as C |
|
|
|
|
| class ConcatMatrix: |
| """Read-only row-concatenation of several memory-mapped (n_i, dim) matrices. |
| |
| Avoids materialising a 21 GB copy just to present S2+S3 as one matrix. |
| Supports the two access patterns blocking uses: contiguous slices and |
| integer-array fancy indexing. |
| """ |
|
|
| def __init__(self, paths): |
| self.parts = [np.load(p, mmap_mode="r") for p in paths] |
| dims = {p.shape[1] for p in self.parts} |
| if len(dims) != 1: |
| raise SystemExit(f"dim mismatch across parts: {dims}") |
| self.offs = np.cumsum([0] + [len(p) for p in self.parts]) |
| self.shape = (int(self.offs[-1]), self.parts[0].shape[1]) |
| self.dtype = self.parts[0].dtype |
|
|
| def __len__(self): |
| return self.shape[0] |
|
|
| def _slice(self, a, b): |
| out = np.empty((b - a, self.shape[1]), dtype=self.dtype) |
| for i, part in enumerate(self.parts): |
| lo, hi = int(self.offs[i]), int(self.offs[i + 1]) |
| s, e = max(a, lo), min(b, hi) |
| if s < e: |
| out[s - a:e - a] = part[s - lo:e - lo] |
| return out |
|
|
| def __getitem__(self, key): |
| if isinstance(key, slice): |
| a, b, _ = key.indices(self.shape[0]) |
| return self._slice(a, b) |
| if isinstance(key, (int, np.integer)): |
| return self._slice(int(key), int(key) + 1)[0] |
| idx = np.asarray(key) |
| out = np.empty((len(idx), self.shape[1]), dtype=self.dtype) |
| for i, part in enumerate(self.parts): |
| lo, hi = int(self.offs[i]), int(self.offs[i + 1]) |
| m = (idx >= lo) & (idx < hi) |
| if m.any(): |
| sub = idx[m] - lo |
| order = np.argsort(sub) |
| rows = part[sub[order]] |
| dest = np.where(m)[0][order] |
| out[dest] = rows |
| return out |
|
|
|
|
| def _check(mat, expect_rows, what): |
| if mat.shape[0] != expect_rows: |
| raise SystemExit(f"{what}: {mat.shape[0]:,} rows but parquet has " |
| f"{expect_rows:,} -- embeddings are not row-aligned") |
|
|
|
|
| def load_emb2(split, n_s1, n_s23): |
| """Second dense channel: the pre-computed e5-large matrices, as (pool, query). |
| |
| No encoding is needed -- the S1 vectors are already on disk in parquet row |
| order, so they are indexed directly by row. |
| """ |
| H = C.HYBRID |
| if split == "train": |
| q = np.load(H["train_s1"], mmap_mode="r") |
| c = ConcatMatrix([H["train_s2"], H["train_s3"]]) |
| elif split == "test": |
| q = np.load(H["test_s1"], mmap_mode="r") |
| c = ConcatMatrix([H["test_s23"]]) |
| else: |
| raise SystemExit(f"unknown split {split!r}") |
| _check(q, n_s1, f"{split} S1 emb2") |
| _check(c, n_s23, f"{split} S2+S3 emb2") |
| return c, q |
|
|
|
|
| def load(split, s1_ids, s23_ids): |
| """-> (pool_matrix, query_matrix, meta) for split in {"train","test"}. |
| |
| s1_ids / s23_ids are the id lists in the order the pipeline built them and |
| are used purely to validate row counts (and, cheaply, the first/last id). |
| """ |
| src = C.EMB.get("source", "internal") |
| if src == "internal": |
| from encoder import load_matrix, read_ids |
| pool_dir, q_dir = C.emb_dir(f"{split}_pool"), C.emb_dir(f"{split}_s1") |
| c_emb, cmeta = load_matrix(pool_dir) |
| q_emb, _ = load_matrix(q_dir) |
| if list(s23_ids) != read_ids(os.path.join(pool_dir, "ids.txt")): |
| raise SystemExit("pool embedding order != S2+S3 order") |
| if list(s1_ids) != read_ids(os.path.join(q_dir, "ids.txt")): |
| raise SystemExit("S1 embedding order != S1 order") |
| return c_emb, q_emb, cmeta |
|
|
| if src != "hybrid_bridge": |
| raise SystemExit(f"unknown EMB['source'] {src!r}") |
|
|
| H = C.HYBRID |
| if split == "train": |
| q_emb = np.load(H["train_s1"], mmap_mode="r") |
| c_emb = ConcatMatrix([H["train_s2"], H["train_s3"]]) |
| elif split == "test": |
| q_emb = np.load(H["test_s1"], mmap_mode="r") |
| c_emb = ConcatMatrix([H["test_s23"]]) |
| else: |
| raise SystemExit(f"unknown split {split!r}") |
|
|
| _check(q_emb, len(s1_ids), f"{split} S1 embeddings") |
| _check(c_emb, len(s23_ids), f"{split} S2+S3 embeddings") |
| meta = {"model": H["model"], "dim": H["dim"], "n": c_emb.shape[0], |
| "source": "hybrid_bridge", "template": H["template"], |
| "query_prefix": H["query_prefix"], "doc_prefix": H["doc_prefix"]} |
| return c_emb, q_emb, meta |
|
|
