File size: 3,083 Bytes
53e24ca f092803 53e24ca f092803 e099c73 | 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 | import math
from types import SimpleNamespace
import pytest
from gemma_rlcd.json_backend import JSONMLXBackend
from gemma_rlcd.json_scoring import FieldTokens
mx = pytest.importorskip("mlx.core")
@pytest.mark.parametrize("batch_size", [1, 2, 3])
def test_complete_candidate_likelihood_uses_suffix_and_ignores_padding(batch_size):
backend = JSONMLXBackend.__new__(JSONMLXBackend)
backend.mx = mx
backend.branch_batch_size = batch_size
backend.tokenizer = SimpleNamespace(pad_token_id=0)
backend.fork_cache = lambda cache, size: []
table = [[-30.0] * 16 for _ in range(16)]
table[9][1], table[9][4] = math.log(0.9), math.log(0.1)
table[1][2], table[1][3] = math.log(0.1), math.log(0.9)
transition = mx.array(table)
backend.model = SimpleNamespace(
language_model=SimpleNamespace(
model=lambda *, inputs, cache, logits_to_keep: inputs[:, -logits_to_keep:, None],
logits_from_hidden=lambda hidden: transition[hidden[..., 0]],
)
)
fields = [(0, FieldTokens((9,), ((1, 2), (1, 3), (4,))))]
streamed = []
results, batches = backend._sequence_scores([], 100, fields, on_scores=streamed.append)
likelihoods = [math.exp(value) for value in results[0].logits]
assert likelihoods == pytest.approx([0.09, 0.81, 0.1], abs=1e-6)
assert results[0].allowed_token_mass == pytest.approx(1, abs=1e-6)
assert sum(batches) == 3
assert max(batches) <= batch_size
assert streamed == [[(0, results[0])]]
def test_repeated_schema_reuses_only_tokenization_and_still_prefills_each_image(
monkeypatch, tmp_path
):
from gemma_rlcd import json_backend
from gemma_rlcd.core import Noul, State, TokenScores
encoded, prefills = [], []
backend = JSONMLXBackend.__new__(JSONMLXBackend)
backend.tokenizer = SimpleNamespace(
encode=lambda text, **kwargs: encoded.append(text) or list(map(ord, text))
)
backend.max_input_tokens = 8192
backend.last_stats = {}
monkeypatch.setattr(
json_backend,
"prepare_generation",
lambda backend, state, questions: (
"prompt",
{"input_ids": SimpleNamespace(shape=(1, 6)), "image": state.images},
),
)
backend.prefill = lambda prepared: prefills.append(prepared.inputs["image"]) or []
backend._sequence_scores = lambda cache, tokens, fields, on_scores: (
{index: TokenScores((float(len(prefills)), 0), 1, 10) for index, _ in fields},
[2],
)
questions = {"visible": Noul("Visible?")}
paths = [tmp_path / name for name in ("one.jpg", "two.jpg")]
for path in paths:
path.touch()
first = backend.score_questions(State(images=(str(paths[0]),)), questions)
assert not backend.last_stats["schema_cache_hit"]
count = len(encoded)
second = backend.score_questions(State(images=(str(paths[1]),)), questions)
assert backend.last_stats["schema_cache_hit"]
assert len(encoded) == count
assert prefills == [(str(paths[0]),), (str(paths[1]),)]
assert first[0].logits != second[0].logits
|