Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 83 additions & 71 deletions programasweights/_runtime_vision.py
Original file line number Diff line number Diff line change
Expand Up @@ -169,14 +169,6 @@ def _eval_prompt(self, parts):

marker = mtmd.mtmd_default_marker().decode("utf-8")
prompt, urls = _to_prompt(*parts, image_marker=marker)
encoded = prompt.encode("utf-8")

input_text = mtmd.mtmd_input_text()
input_text.text = encoded
input_text.text_len = len(encoded)
input_text.add_special = False
input_text.parse_special = True

bitmaps = []
chunks = None
try:
Expand All @@ -188,70 +180,90 @@ def _eval_prompt(self, parts):
raise RuntimeError("Failed to create image bitmap.")
bitmaps.append(bitmap)

chunks = mtmd.mtmd_input_chunks_init()
if not chunks:
raise RuntimeError("Failed to allocate prompt chunks.")

bitmap_array = (
mtmd.mtmd_bitmap_p_ctypes * len(bitmaps)
)(*bitmaps)
result = mtmd.mtmd_tokenize(
handler.mtmd_ctx,
chunks,
ctypes.byref(input_text),
bitmap_array,
len(bitmaps),
)
if result != 0:
raise ValueError(f"Prompt tokenization failed: {result}.")

for index in range(mtmd.mtmd_input_chunks_size(chunks)):
chunk = mtmd.mtmd_input_chunks_get(chunks, index)
if not chunk:
raise RuntimeError("Missing prompt chunk.")

size = mtmd.mtmd_input_chunk_get_n_tokens(chunk)
if model.n_tokens + size >= model.n_ctx():
raise ValueError("Prompt leaves no room for output.")

kind = mtmd.mtmd_input_chunk_get_type(chunk)
if kind == mtmd.MTMD_INPUT_CHUNK_TYPE_TEXT:
count = ctypes.c_size_t()
tokens = mtmd.mtmd_input_chunk_get_tokens_text(
chunk, ctypes.byref(count)
)
if count.value:
if not tokens:
raise RuntimeError("Missing text tokens.")
model.eval([tokens[i] for i in range(count.value)])

elif kind == mtmd.MTMD_INPUT_CHUNK_TYPE_IMAGE:
start = model.n_tokens
end = self.lib.llama_pos(0)
result = mtmd.mtmd_helper_eval_chunk_single(
handler.mtmd_ctx,
model._ctx.ctx,
chunk,
self.lib.llama_pos(start),
self.lib.llama_seq_id(0),
model.n_batch,
True,
ctypes.byref(end),
)
if result != 0:
raise RuntimeError(
f"Image evaluation failed: {result}."
# PAW images are independent stills. Some native Qwen backends
# merge adjacent same-size bitmaps as video frames, changing the
# trained representation. Tokenize one image per call without
# inserting separators or changing any literal prompt text.
text_segments = prompt.split(marker)
groups = [(prompt, [])] if not bitmaps else [
(text_segments[i] + marker +
(text_segments[-1] if i == len(bitmaps) - 1 else ""), [bitmap])
for i, bitmap in enumerate(bitmaps)
]
for text, group in groups:
encoded = text.encode("utf-8")
input_text = mtmd.mtmd_input_text()
input_text.text = encoded
input_text.text_len = len(encoded)
input_text.add_special = False
input_text.parse_special = True
chunks = mtmd.mtmd_input_chunks_init()
if not chunks:
raise RuntimeError("Failed to allocate prompt chunks.")

bitmap_array = (
mtmd.mtmd_bitmap_p_ctypes * len(group)
)(*group)
result = mtmd.mtmd_tokenize(
handler.mtmd_ctx,
chunks,
ctypes.byref(input_text),
bitmap_array,
len(group),
)
if result != 0:
raise ValueError(f"Prompt tokenization failed: {result}.")

for index in range(mtmd.mtmd_input_chunks_size(chunks)):
chunk = mtmd.mtmd_input_chunks_get(chunks, index)
if not chunk:
raise RuntimeError("Missing prompt chunk.")

size = mtmd.mtmd_input_chunk_get_n_tokens(chunk)
if model.n_tokens + size >= model.n_ctx():
raise ValueError("Prompt leaves no room for output.")

kind = mtmd.mtmd_input_chunk_get_type(chunk)
if kind == mtmd.MTMD_INPUT_CHUNK_TYPE_TEXT:
count = ctypes.c_size_t()
tokens = mtmd.mtmd_input_chunk_get_tokens_text(
chunk, ctypes.byref(count)
)
if count.value:
if not tokens:
raise RuntimeError("Missing text tokens.")
model.eval([tokens[i] for i in range(count.value)])

elif kind == mtmd.MTMD_INPUT_CHUNK_TYPE_IMAGE:
start = model.n_tokens
end = self.lib.llama_pos(0)
result = mtmd.mtmd_helper_eval_chunk_single(
handler.mtmd_ctx,
model._ctx.ctx,
chunk,
self.lib.llama_pos(start),
self.lib.llama_seq_id(0),
model.n_batch,
True,
ctypes.byref(end),
)
if not start < end.value < model.n_ctx():
raise RuntimeError("Invalid image end position.")

# Image embeddings occupy these positions. These zeros
# are Python bookkeeping, not tokens sent to the model.
model.input_ids[start:end.value] = 0
model.n_tokens = end.value
model._requires_eval = False
else:
raise ValueError("Unsupported prompt chunk type.")
if result != 0:
raise RuntimeError(
f"Image evaluation failed: {result}."
)
if not start < end.value < model.n_ctx():
raise RuntimeError("Invalid image end position.")

# Image embeddings occupy these positions. These zeros
# are Python bookkeeping, not tokens sent to the model.
model.input_ids[start:end.value] = 0
model.n_tokens = end.value
model._requires_eval = False
else:
raise ValueError("Unsupported prompt chunk type.")

mtmd.mtmd_input_chunks_free(chunks)
chunks = None

if model.n_tokens == 0:
raise ValueError("The prompt produced no tokens.")
Expand Down
21 changes: 19 additions & 2 deletions tests/test_vision_prompt.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def tolist(self):
@pytest.fixture
def native():
state = NS(chunks=[], flags=None, loaded=[], freed=[], chunk_freed=False,
images=[], fail_tokenize=False, fail_image=False)
images=[], fail_tokenize=False, fail_image=False, prompts=[])
model = NS(n_tokens=0, input_ids=Vector([0] * 256), n_ctx=lambda: 256,
_ctx=NS(ctx=None), n_batch=128)
def evaluate(tokens):
Expand All @@ -54,6 +54,7 @@ def tokenize(ctx, chunks, ptr, bitmaps, count):
prompt = ptr._obj
state.flags = (prompt.add_special, prompt.parse_special)
state.text = prompt.text[:prompt.text_len].decode()
state.prompts.append(state.text)
state.chunks = []
for i, piece in enumerate(state.text.split("<__media__>")):
if i:
Expand Down Expand Up @@ -103,7 +104,8 @@ def test_exact_text_and_image_order_no_wrapping_or_special_token_insertion(nativ
state, handler, model, lib = native
prompt = evaluate_prompt(handler, model, [text(" pre"), text("fix "), image("B"),
text("mid"), image("A"), image("B"), text("end\n ")], lib)
assert state.text == " prefix <__media__>mid<__media__><__media__>end\n "
assert "".join(state.prompts) == " prefix <__media__>mid<__media__><__media__>end\n "
assert all(p.count("<__media__>") <= 1 for p in state.prompts)
assert state.flags == (False, True)
assert state.loaded == ["B", "A", "B"]
assert state.images == [(1, True), (2, True), (3, True)]
Expand Down Expand Up @@ -266,3 +268,18 @@ def test_evaluation_error_prevents_generation():
with pytest.raises(RuntimeError, match='native failure'):
_Runtime._complete_prompt(owner, ('input',), max_tokens=8, temperature=0)
owner.llm.create_completion.assert_not_called()


def test_adjacent_images_are_not_exposed_to_native_video_merging(native):
state, handler, model, lib = native
tokenize = handler._mtmd_cpp.mtmd_tokenize
def merges_video(*args):
result = tokenize(*args)
if "<__media__><__media__>" in state.text:
state.chunks = [c for i, c in enumerate(state.chunks)
if not (i and c.kind == state.chunks[i - 1].kind == 2)]
return result
handler._mtmd_cpp.mtmd_tokenize = merges_video
evaluate_prompt(handler, model, [image("A"), image("B")], lib)
assert state.images == [(1, True), (2, True)]
assert "".join(state.prompts) == "<__media__><__media__>"
Loading