File size: 6,249 Bytes
e79d917 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """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) # memmap gather wants sorted reads
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
|