File size: 3,661 Bytes
38d2d22 | 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 | --- a/python/src/captchakraken/pointer/service.py
+++ b/python/src/captchakraken/pointer/service.py
@@ -29,6 +29,8 @@
from . import vllm_plugin
+VLLM_VIDEO = "<|vision_start|><|video_pad|><|vision_end|>"
+
class PointerService:
def __init__(self, checkpoint: Union[str, Path], detector: Union[str, Path, None] = None,
gpu_memory_utilization: float = 0.45, max_model_len: int = 8192, **llm_kwargs):
@@ -41,7 +43,10 @@
self.llm = LLM(model=str(checkpoint), runner="pooling", max_model_len=max_model_len,
gpu_memory_utilization=gpu_memory_utilization,
limit_mm_per_prompt={"image": 1 + L.MAX_HISTORY, "video": 1},
- mm_processor_kwargs={"min_pixels": 64 * 64, "max_pixels": L.MAX_PIXELS}, **llm_kwargs)
+ mm_processor_kwargs={"min_pixels": 64 * 64, "max_pixels": L.MAX_PIXELS},
+ # The heads read every token's hidden state; vLLM returns none for a
+ # prefix it served from cache, so a repeated prompt came back short.
+ enable_prefix_caching=False, **llm_kwargs)
self.detector = None
if detector is not None:
from .detector import Detector
@@ -59,7 +64,7 @@
marked: bool = True, history: Sequence[Image.Image] = ()) -> Dict:
"""The request vLLM gets, and what decoding its answer needs. `history`:
the puzzle's earlier screens, oldest first, as they were sent (marked the
- same way), so the prompt up to them is the one already cached."""
+ same way). The prefix is recomputed each time: the heads need its hidden states."""
if (image is None) == (frames is None):
raise ValueError("give exactly one of image or frames")
if frames is not None and history:
@@ -78,7 +83,12 @@
if tracks:
fr = np.stack([np.asarray(M.draw(im, E.on_frame(tracks, k)))
for k, im in enumerate(ims)])
- return {"request": {"prompt": L.chat(self.processor, True),
+ # vLLM (0.25.1) expands a video only where it finds the template's whole
+ # <|vision_start|><|video_pad|><|vision_end|>; `chat` sends transformers
+ # the bare pad. Both then give the same tokens but one: training merges the
+ # instruction's last "." with the first timestamp's "<", vLLM splits them.
+ prompt = L.chat(self.processor, True).replace(L.VIDEO_PAD, VLLM_VIDEO)
+ return {"request": {"prompt": prompt,
"multi_modal_data": {"video": (fr, meta)},
"mm_processor_kwargs": {"do_sample_frames": False}},
"size": (w, h), "elements": tracks, "shown": None,
--- a/python/src/captchakraken/pointer/vllm_plugin.py
+++ b/python/src/captchakraken/pointer/vllm_plugin.py
@@ -65,6 +65,9 @@
# A real prompt without its markers still fails in Positions.
out.append(torch.zeros(L.PACKED + TRAILER, device=h.device))
continue
+ if h.shape[0] != ids.shape[0]:
+ raise RuntimeError(f"{h.shape[0]} hidden states for a {ids.shape[0]}-token prompt: "
+ "the heads need them all (is prefix caching on?)")
pos = Positions(ids, m.pointer_token_ids)
v = m.pointer_heads.decide(h, pos, m.language_model.lm_head.weight)
out.append(torch.cat([v, torch.tensor([float(pos.per_moment), float(pos.n_moments)],
|