| import json |
|
|
| import pytest |
|
|
| from groot1.schema import SystemOneRequest, Usage |
| from groot1.serve import GrootEngine |
|
|
|
|
| class FakeMember: |
| def __init__(self, probs, p_correct, cfg=None): |
| self.probs, self.p_correct, self.cfg = probs, p_correct, cfg |
|
|
| def forward(self, state, questions): |
| extra = [{"p_correct": self.p_correct, "truncated": False} for _ in questions] |
| return self.probs, extra, Usage(input_tokens=10, output_tokens=len(questions)) |
|
|
|
|
| def ensemble(members, weights): |
| eng = object.__new__(GrootEngine) |
| eng.members, eng.weights, eng.name = members, weights, "groot-1" |
| return eng |
|
|
|
|
| def test_ensemble_averages_member_probabilities(): |
| req = SystemOneRequest(state="x", questions={ |
| "c": {"type": "choice", "instructions": "pick", "criteria": {"a": None, "b": None}}, |
| "n": {"type": "noul", "instructions": "true?"}, |
| }) |
| eng = ensemble([FakeMember([[0.9, 0.1], [0.2, 0.8]], 0.8), FakeMember([[0.3, 0.7], [0.6, 0.4]], 0.4)], [0.5, 0.5]) |
| res = eng.answer(req.state, req.questions, extended=True) |
| assert res.answers["c"].probabilities == pytest.approx({"a": 0.6, "b": 0.4}) |
| assert res.answers["c"].choice == "a" |
| assert res.answers["n"].noul == pytest.approx(0.6) |
| assert res.x_groot["c"]["p_correct"] == pytest.approx(0.6) |
| assert res.model == "groot-1" |
|
|
|
|
| def test_ensemble_weights_are_respected(): |
| req = SystemOneRequest(state="x", questions={"c": {"type": "choice", "instructions": "pick", "criteria": {"a": None, "b": None}}}) |
| eng = ensemble([FakeMember([[1.0, 0.0]], 1.0), FakeMember([[0.0, 1.0]], 0.0)], [0.25, 0.75]) |
| assert eng.answer(req.state, req.questions).answers["c"].probabilities == pytest.approx({"a": 0.25, "b": 0.75}) |
|
|
|
|
| def test_ensemble_spec_loads_members(tmp_path, monkeypatch): |
| (tmp_path / "ensemble.json").write_text(json.dumps({"name": "groot-1", "members": ["s", "m"], "weights": [1, 3]})) |
| loaded = [] |
| real_init = GrootEngine.__init__ |
|
|
| def fake_init(self, checkpoint, device=None, model_name=None, overrides=None): |
| if (tmp_path / "ensemble.json") == (type(tmp_path)(checkpoint) / "ensemble.json"): |
| return real_init(self, checkpoint, "cpu", model_name, overrides) |
| loaded.append(type(tmp_path)(checkpoint).name) |
| self.members, self.cfg = [], object() |
|
|
| monkeypatch.setattr(GrootEngine, "__init__", fake_init) |
| eng = GrootEngine(tmp_path) |
| assert loaded == ["s", "m"] |
| assert eng.weights == pytest.approx([0.25, 0.75]) |
| assert eng.name == "groot-1" and eng.aliases["jev-latest"] == "groot-1" |
|
|
