diff --git a/conditioning.py b/conditioning.py index 95721795..16c93a5e 100644 --- a/conditioning.py +++ b/conditioning.py @@ -5162,20 +5162,93 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): return False return True - def _v2_clip_similarity_scores(self, clip, phrase, category_vectors, encode_cache=None): - phrase_cond, _, _ = self._v2_encode_prompt(clip, phrase, encode_cache=encode_cache, - purpose="phrase classification") - phrase_vector = self._v2_conditioning_vector(phrase_cond) + def _v2_phrase_vectors_in_one_pass(self, clip, texts, encode_cache=None): + """{phrase: vector} for many phrases from ONE encode, or None if that cannot be done. + + Classifying a phrase means comparing its vector to eight category vectors. Getting + that vector by encoding the phrase on its own costs a full pass of the text encoder + EACH — on H3 that is Qwen3-VL-32B, once per phrase the heuristic was unsure about. + + The phrases are encoded together instead, as one comma-joined text, and each phrase's + vector is the mean of the rows its own tokens occupy. The tokenizer's offset mapping + says which rows those are exactly, so this is a lookup rather than a guess. A phrase + also gets a little of its neighbours' context this way, which is closer to how it + appears in a real prompt than the phrase standing alone. + + Returns None — caller falls back to one encode per phrase — when the tokenizer offers + no offsets, or when the spans do not line up with the conditioning that came back. + Nothing here is worth a wrong classification. + """ + texts = [t for t in dict.fromkeys(str(t or "").strip() for t in texts) if t] + if len(texts) < 2: + return None + try: + from . import h3_token_weights as _tw + except ImportError: + import h3_token_weights as _tw + tokenizer = _tw.h3_tokenizer(clip) + if tokenizer is None: + return None + + joined = ", ".join(texts) + cond, meta, _ = self._v2_encode_prompt(clip, joined, encode_cache=encode_cache, + purpose="phrase classification (batched)") + if not isinstance(cond, torch.Tensor) or cond.dim() < 3: + return None + try: + offsets = list(tokenizer(joined, add_special_tokens=False, + return_offsets_mapping=True)["offset_mapping"]) + except Exception: # noqa: BLE001 — no offsets, no batching; fall back to one each + return None + prompt_tokens = len(offsets) + if prompt_tokens <= 0: + return None + cond_len = int(cond.shape[1]) + base = _tw.prompt_base((meta or {}).get("minimax_token_tags"), cond_len, prompt_tokens) + if base is None: + return None + + # Each phrase is located at the offset it was JOINED at, not by searching the text: + # one phrase can be a substring of another ("dress" inside "red dress"), and a search + # would hand both the same rows. The join is ours, so the positions are known. + rows = cond[0] + out = {} + cursor = 0 + for text in texts: + c_start, c_end = cursor, cursor + len(text) + cursor = c_end + 2 # the ", " this was joined with + span = _tw.token_spans_from_offsets(offsets, [(c_start, c_end, 1.0)]) + if not span: + return None + a, b = base + int(span[0][0]), base + int(span[0][1]) + if a < 0 or b > cond_len or b <= a: + return None + out[text] = rows[a:b].float().mean(dim=0) + return out if len(out) == len(texts) else None + + def _v2_scores_against_categories(self, phrase_vector, category_vectors): + """Cosine similarity of one phrase vector against each category vector, 0..1. + + Shared by both ways of getting that vector — the batched span pool and the per-phrase + encode — so a phrase scores the same however its vector was obtained. + """ if phrase_vector is None: return {} scores = {} - for category, vector in category_vectors.items(): + for category, vector in (category_vectors or {}).items(): if not isinstance(vector, torch.Tensor) or vector.shape != phrase_vector.shape: continue - sim = F.cosine_similarity(phrase_vector.unsqueeze(0), vector.to(phrase_vector.device).unsqueeze(0), dim=-1) + sim = F.cosine_similarity(phrase_vector.unsqueeze(0), + vector.to(phrase_vector.device).unsqueeze(0), dim=-1) scores[category] = float(((sim.item() + 1.0) * 0.5)) return scores + def _v2_clip_similarity_scores(self, clip, phrase, category_vectors, encode_cache=None): + phrase_cond, _, _ = self._v2_encode_prompt(clip, phrase, encode_cache=encode_cache, + purpose="phrase classification") + return self._v2_scores_against_categories( + self._v2_conditioning_vector(phrase_cond), category_vectors) + # Eight fixed strings embedded to classify phrases by cosine similarity. They are # CONSTANT — CATEGORY_DESCRIPTIONS is a class attribute, so the result depends only on the # encoder. Held for the process rather than the run, keyed by the encoder object, because @@ -5664,8 +5737,18 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): if uncertain and clip is not None: category_vectors = self._v2_category_vectors(clip, encode_cache=encode_cache) + # One encode for every uncertain phrase, rather than one EACH. None means the + # batch could not be trusted (no offsets, or spans that did not line up), and the + # per-phrase path below runs unchanged. + batched = self._v2_phrase_vectors_in_one_pass( + clip, [i["text"] for i in uncertain], encode_cache=encode_cache) for item in uncertain: - clip_scores = self._v2_clip_similarity_scores(clip, item["text"], category_vectors, encode_cache=encode_cache) + vector = (batched or {}).get(item["text"]) + clip_scores = ( + self._v2_scores_against_categories(vector, category_vectors) + if vector is not None else + self._v2_clip_similarity_scores(clip, item["text"], category_vectors, encode_cache=encode_cache) + ) if not clip_scores: continue merged = self._v2_merge_clip_category_scores(item["category_scores"], clip_scores) diff --git a/tests/test_batched_phrase_vectors.py b/tests/test_batched_phrase_vectors.py new file mode 100644 index 00000000..2359e0a2 --- /dev/null +++ b/tests/test_batched_phrase_vectors.py @@ -0,0 +1,95 @@ +"""Classifying N uncertain phrases costs one encode, not N. + +A phrase's vector used to come from encoding that phrase alone — a full pass of the text +encoder each, and on H3 that encoder is Qwen3-VL-32B. The phrases are encoded together now +and each one's vector is the mean of the rows its own tokens occupy, which the tokenizer's +offset mapping identifies exactly. +""" +import sys +import types + +import pytest +import torch + +sys.path.insert(0, ".") + + +class _Tok: + """Offset mapping over whitespace/comma-free spans, like a real subword tokenizer.""" + def __call__(self, text, add_special_tokens=False, return_offsets_mapping=False): + offsets, i = [], 0 + for chunk in text.split(" "): + if chunk: + offsets.append((i, i + len(chunk))) + i += len(chunk) + 1 + return {"offset_mapping": offsets} + + +@pytest.fixture +def refiner(monkeypatch): + import conditioning as C + r = C.FunPackVideoRefinerV2() + r._v2_reset_encode_tally() + return r + + +@pytest.fixture +def wired(monkeypatch, refiner): + """One encode returns a conditioning with one row per token of the joined text.""" + import conditioning as C + import h3_token_weights as tw + calls = {"n": 0} + + def fake_encode(self, clip, text, encode_cache=None, purpose="prompt", **kw): + calls["n"] += 1 + n = len(_Tok()(text)["offset_mapping"]) + rows = torch.arange(n, dtype=torch.float32).unsqueeze(-1).repeat(1, 4) + return rows.unsqueeze(0), {"minimax_token_tags": [1] * n}, "" + + monkeypatch.setattr(C.FunPackVideoRefinerV2, "_v2_encode_prompt", fake_encode) + monkeypatch.setattr(tw, "h3_tokenizer", lambda clip: _Tok()) + return refiner, calls + + +def test_many_phrases_cost_one_encode(wired): + refiner, calls = wired + out = refiner._v2_phrase_vectors_in_one_pass(object(), ["alpha", "beta", "gamma"]) + assert out is not None and set(out) == {"alpha", "beta", "gamma"} + assert calls["n"] == 1 + + +def test_each_phrase_gets_its_own_rows(wired): + """The rows are numbered, so the pooled value says which span was read.""" + refiner, _ = wired + out = refiner._v2_phrase_vectors_in_one_pass(object(), ["alpha", "beta"]) + # joined = "alpha, beta" -> tokens ["alpha,", "beta"] at rows 0 and 1 + assert float(out["alpha"][0]) == 0.0 + assert float(out["beta"][0]) == 1.0 + + +def test_a_phrase_inside_another_is_not_given_its_rows(wired): + """'dress' is a substring of 'red dress'. Searching the text would hand both the same + rows; the positions come from the join instead.""" + refiner, _ = wired + out = refiner._v2_phrase_vectors_in_one_pass(object(), ["red dress", "dress"]) + assert not torch.equal(out["red dress"], out["dress"]) + + +def test_one_phrase_is_not_worth_batching(wired): + refiner, calls = wired + assert refiner._v2_phrase_vectors_in_one_pass(object(), ["alpha"]) is None + assert calls["n"] == 0 + + +def test_no_offsets_means_fall_back(monkeypatch, refiner): + import h3_token_weights as tw + monkeypatch.setattr(tw, "h3_tokenizer", lambda clip: None) + assert refiner._v2_phrase_vectors_in_one_pass(object(), ["a", "b"]) is None + + +def test_both_paths_score_a_vector_the_same_way(refiner): + """The cosine step is shared, so a phrase scores the same however its vector was got.""" + cats = {"action": torch.tensor([1.0, 0.0]), "style": torch.tensor([0.0, 1.0])} + scores = refiner._v2_scores_against_categories(torch.tensor([1.0, 0.0]), cats) + assert scores["action"] > scores["style"] + assert refiner._v2_scores_against_categories(None, cats) == {}