From b0c22c85ddf3ba8cdcc0aa8aa60dd837a927869d Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Tue, 25 Aug 2026 13:36:23 +0300 Subject: [PATCH] perf(studio): classify every uncertain phrase in one encode instead of one each MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A phrase's vector came from encoding that phrase alone — a full pass of the text encoder per phrase, and on H3 that encoder is Qwen3-VL-32B. The phrases are encoded together now, as one joined text, and each vector is the mean of the rows its own tokens occupy; the tokenizer's offset mapping identifies those rows exactly, so this is a lookup rather than an approximation. Each phrase is located at the position it was joined at rather than by searching the text, because one phrase can be a substring of another and a search would hand both the same rows. Falls back to one encode per phrase, unchanged, when the tokenizer offers no offsets or the spans do not line up with the conditioning that came back. The cosine step is shared by both paths so a phrase scores the same either way, and the batched pass is labelled separately in the per-run encoder tally. Co-Authored-By: Claude Opus 5 --- conditioning.py | 97 ++++++++++++++++++++++++++-- tests/test_batched_phrase_vectors.py | 95 +++++++++++++++++++++++++++ 2 files changed, 185 insertions(+), 7 deletions(-) create mode 100644 tests/test_batched_phrase_vectors.py 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) == {}