File size: 6,339 Bytes
53e24ca f092803 53e24ca f092803 e099c73 f092803 e099c73 f092803 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | import json
import pytest
from gemma_rlcd.comparison import discrete_answers, generation_task, parse_generated
from gemma_rlcd.core import Choice, Independent, Noul, Score, State
def questions():
return {
"brand": Choice("Which brand is this?", {"Porsche": "Porsche", "Mercedes": "Mercedes"}),
"quality": Score("Grade its condition", ["Poor", "Good", "Excellent"]),
"moving": Noul("Is it moving?"),
"presence": Independent("Which are present?", {"car": "A car", "person": "A person"}),
}
def test_generation_requests_all_fields_in_one_json_answer():
task = generation_task(questions())
assert "Porsche" in task and "Mercedes" in task
assert "Grade its condition" in task and "Excellent" in task
assert "a boolean" in task and "proposition name to a boolean" in task
expected = {
"brand": "Porsche",
"quality": 2,
"moving": False,
"presence": {"car": True, "person": False},
}
assert parse_generated(json.dumps(expected), questions()) == expected
assert parse_generated("```json\n" + json.dumps(expected) + "\n```", questions()) == expected
@pytest.mark.parametrize(
"text",
[
'{"brand":"BMW","quality":2,"moving":false,"presence":{"car":true,"person":false}}',
'{"brand":"Porsche","quality":true,"moving":false,"presence":{"car":true,"person":false}}',
'{"brand":"Porsche","quality":3,"moving":false,"presence":{"car":true,"person":false}}',
'{"brand":"Porsche","quality":2,"moving":false,"presence":{"car":true}}',
'{"brand":"Porsche","brand":"Mercedes"}',
'The answer is {"brand":"Porsche"}',
],
)
def test_generated_output_must_match_the_contract(text):
with pytest.raises(ValueError):
parse_generated(text, questions())
def test_agreement_uses_grade_mode_and_boolean_thresholds():
assert discrete_answers(
{
"grade": {
"type": "score",
"score": 1.4,
"probabilities": {"0": 0.1, "1": 0.4, "2": 0.5},
},
"moving": {"type": "noul", "noul": 0.7},
"presence": {"type": "independent", "probabilities": {"car": 0.9, "person": 0.1}},
}
) == {"grade": 2, "moving": True, "presence": {"car": True, "person": False}}
def test_streaming_generation_uses_real_chunks_and_validates_complete_answer(monkeypatch):
import sys
from types import SimpleNamespace
from gemma_rlcd import comparison
seen = []
parts = ['{"visible":', "true", "}"]
def stream(*args, **kwargs):
for index, text in enumerate(parts):
assert len(seen) == index
yield SimpleNamespace(
text=text,
generation_tokens=index + 1,
prompt_tokens=50,
finish_reason="stop" if index == len(parts) - 1 else None,
)
monkeypatch.setitem(
sys.modules, "mlx_vlm", SimpleNamespace(generate=None, stream_generate=stream)
)
monkeypatch.setattr(comparison, "prepare_generation", lambda *args: ("prompt", {}))
backend = SimpleNamespace(
model=None,
processor=None,
tokenizer=SimpleNamespace(encode=lambda *args, **kwargs: [1]),
mx=SimpleNamespace(synchronize=lambda: None),
)
result = comparison.generate_answers(
backend,
State(text="test"),
{"visible": Noul("Visible?")},
on_token=lambda text, count: seen.append((text, count)),
)
assert seen == list(zip(parts, [1, 2, 3], strict=True))
assert result["valid"] and result["answers"] == {"visible": True}
assert result["raw_text"] == "".join(parts)
def test_concurrent_comparison_streams_both_paths_before_either_finishes(monkeypatch):
from threading import Event
from types import SimpleNamespace
from gemma_rlcd import comparison
from gemma_rlcd.core import TokenScores
scored, generated = Event(), Event()
events = []
class Backend:
last_stats = {}
processor = SimpleNamespace(tokenizer=SimpleNamespace(mutable=[]))
def symbols(self, count):
return ("A", "B")
def score_questions(self, state, questions, on_scores, on_progress=None):
assert generated.wait(2), "Generation never started alongside scoring"
result = TokenScores((5, 0), 1, 20)
on_scores([(0, result)])
scored.set()
return [result]
backend = Backend()
def generate(worker, state, questions, on_token, on_progress=None):
assert worker.processor is not backend.processor
assert worker.processor.tokenizer is not backend.processor.tokenizer
on_token('{"visible":', 1)
generated.set()
assert scored.wait(2), "Scoring never completed while generation was active"
on_token("true}", 2)
return {"answers": {"visible": True}, "valid": True, "inference_seconds": 0}
monkeypatch.setattr(comparison, "generate_answers", generate)
_, result = comparison.compare(
backend,
State(text="scene"),
{"visible": Noul("Visible?")},
0,
emit=events.append,
concurrent=True,
)
assert events[0]["type"] == "race_start"
first_token = next(i for i, event in enumerate(events) if event["type"] == "token")
first_answer = next(i for i, event in enumerate(events) if event["type"] == "answer")
first_finish = next(i for i, event in enumerate(events) if event["type"] == "phase_complete")
assert first_token < first_answer < first_finish
assert result["agreement"] == {"visible": True}
assert result["methodology"]["execution"] == "concurrent_shared_gpu"
@pytest.mark.parametrize("cancel_at", ["race_start", "phase_start"])
def test_concurrent_comparison_can_stop_before_or_after_worker_launch(cancel_at):
from types import SimpleNamespace
from gemma_rlcd.comparison import compare
def emit(event):
if event["type"] == cancel_at:
raise RuntimeError("Client stopped")
with pytest.raises(RuntimeError, match="Client stopped"):
compare(
SimpleNamespace(),
State(text="scene"),
{"visible": Noul("Visible?")},
0,
emit=emit,
concurrent=True,
)
|