Spaces:
Running
Running
File size: 3,578 Bytes
a1dd5ba | 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 | """The `pe` channel: Meta Perception Encoder (published SOTA zero-shot
video-text retrieval), adopted by measurement — best single-encoder
precision@10 on the bench (48/140 vs shipping fusion 40), complementary
profile (put-into-drawer 10/10 where others fail). Naive equal fusion
measured WORSE than PE alone, so this channel routes through the learned
per-query-type weights like every other.
8 frame vectors per recording in `pe_vectors`; recording score =
top-5-mean of frame cosines (the winning pooling from the zero-shot
sweep). Query cost: one open_clip text forward (~15 ms, cached)."""
from __future__ import annotations
import numpy as np
_TEXT = {}
_IDX = {}
def _text_vec(text):
import torch
if text in _TEXT.get("cache", {}):
return _TEXT["cache"][text]
if "model" not in _TEXT:
import open_clip
from .device import pick, strip_vision, text_only
dev, dtype = pick()
m, _, _ = open_clip.create_model_and_transforms(
"PE-Core-L-14-336", pretrained="meta")
m = strip_vision(m.to(dev).eval(), "visual")
if text_only():
# serving keeps only the text tower; at reduced precision
# it is a few hundred MB instead of the fp32 gigabytes
m = m.to(dtype)
_TEXT["model"] = m
_TEXT["dtype"] = m.ln_final.weight.dtype \
if hasattr(m, "ln_final") else None
_TEXT["tok"] = open_clip.get_tokenizer("PE-Core-L-14-336")
_TEXT["dev"] = dev
_TEXT["cache"] = {}
with torch.no_grad():
t = _TEXT["model"].encode_text(
_TEXT["tok"]([text]).to(_TEXT["dev"]))
v = (t / t.norm(dim=-1, keepdim=True))[0].cpu().float().numpy()
if len(_TEXT["cache"]) > 256:
_TEXT["cache"].clear()
_TEXT["cache"][text] = v
return v
def pe_lookup(store, text):
"""(lookup(stream, t0, t1) -> top5-mean cosine | nan, candidates)."""
from .embeddings import _vec_table
ver = store.table("pe_vectors").state().version
key = (str(store.dir), ver)
if key not in _IDX:
tbl, _ = _vec_table(store, "pe_vectors")
ss = np.asarray(tbl.column("stream").to_pylist())
sa = np.asarray([int(v) for v in tbl.column("ts").to_pylist()])
sb = np.asarray([int(v) for v in tbl.column("t1").to_pylist()])
# rows grouped by recording (stream, t0, t1): build run map
recs = {}
order = np.lexsort((sa, ss))
for r in order:
recs.setdefault((str(ss[r]), int(sa[r]), int(sb[r])),
[]).append(int(r))
idx = {}
for (s, a, b), rows in recs.items():
idx.setdefault(s, []).append((a, b, np.array(rows)))
for s in idx:
idx[s].sort(key=lambda x: x[0])
if len(_IDX) > 8:
_IDX.clear()
_IDX[key] = idx
idx = _IDX[key]
_, vecs = _vec_table(store, "pe_vectors")
sc = vecs @ _text_vec(text)
def rec_score(rows):
v = sc[rows]
k = min(5, len(v))
return float(np.sort(v)[-k:].mean())
def lookup(s, a, b):
lst = idx.get(s)
if not lst:
return float("nan")
starts = [x[0] for x in lst]
j = int(np.searchsorted(starts, a, side="right")) - 1
if j >= 0 and b <= lst[j][1] + 1:
return rec_score(lst[j][2])
return float("nan")
cands = []
for s, lst in idx.items():
for a, b, rows in lst:
cands.append((s, a, b, rec_score(rows)))
cands.sort(key=lambda x: -x[3])
return lookup, cands[:64]
|