Skip to content

Commit c86251a

Browse files
fix(models): report per-file checksum pins in OnnxEncoder.model_sha256 (#403)
#346 / PR #396 added ModelConfig.file_sha256 so every model file can be pinned, and registry.pinned_checksums(cfg) merges it with the primary-file sha256 pin. verify() and pull() use the merged pins, but OnnxEncoder.model_sha256 still read only cfg.sha256. A model whose primary file is pinned only through file_sha256 was reported as "unverified", and that value also keys the embedding cache and the model evidence metadata. model_sha256 now reports pinned_checksums(cfg).get(cfg.files[0]), falling back to "unverified". The other cfg.sha256 readers were checked and are already correct: pinned_checksums() is the merge itself, and the status() note tests "cfg.sha256 or cfg.file_sha256". Behaviour for the shipped registry, where nothing is pinned yet, is unchanged.
1 parent cbd1967 commit c86251a

2 files changed

Lines changed: 61 additions & 2 deletions

File tree

‎src/freshdata/models/runtime.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,14 @@
1717
import numpy as np
1818

1919
from ._lazy import has_semantic_extra, require_onnxruntime, require_tokenizers
20-
from .registry import COL_ENCODER_ID, get_config, is_installed, model_dir, verify
20+
from .registry import (
21+
COL_ENCODER_ID,
22+
get_config,
23+
is_installed,
24+
model_dir,
25+
pinned_checksums,
26+
verify,
27+
)
2128
from .types import ModelError
2229

2330
_STUB_ENV = "FRESHDATA_STUB_ENCODER"
@@ -47,7 +54,9 @@ class OnnxEncoder:
4754
def __init__(self, model_id: str) -> None:
4855
cfg = get_config(model_id)
4956
self.model_id = model_id
50-
self.model_sha256 = cfg.sha256 or "unverified"
57+
# The primary file's pin, whether set via ``sha256`` or ``file_sha256``.
58+
primary_pin = pinned_checksums(cfg).get(cfg.files[0]) if cfg.files else None
59+
self.model_sha256 = primary_pin or "unverified"
5160
self.dim = 384
5261
self._session: Any = None
5362
self._tokenizer: Any = None

‎tests/test_models_runtime.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,14 @@
22

33
from __future__ import annotations
44

5+
import dataclasses
6+
import hashlib
7+
58
import numpy as np
69
import pytest
710

811
from freshdata.models import ModelError, runtime
12+
from freshdata.models import registry as reg
913
from freshdata.models.stub import StubEncoder
1014
from freshdata.semantic.cache import EmbeddingCache
1115

@@ -123,6 +127,52 @@ def test_embedding_cache_disabled():
123127
assert len(cache) == 0
124128

125129

130+
_FAKE_MODEL = b"FAKE-ONNX-BYTES"
131+
_FAKE_TOKENIZER = b'{"vocab": []}'
132+
133+
134+
def _sha(payload: bytes) -> str:
135+
return hashlib.sha256(payload).hexdigest()
136+
137+
138+
def test_onnx_encoder_reports_primary_file_pin_from_file_sha256(monkeypatch):
139+
# Constructing OnnxEncoder is lazy: no ONNX runtime or model files needed.
140+
cfg = dataclasses.replace(
141+
reg.get_config(reg.COL_ENCODER_ID),
142+
sha256=None,
143+
file_sha256=(
144+
("tokenizer.json", _sha(_FAKE_TOKENIZER)),
145+
("model.onnx", _sha(_FAKE_MODEL)),
146+
),
147+
)
148+
monkeypatch.setitem(reg.REGISTRY, reg.COL_ENCODER_ID, cfg)
149+
encoder = runtime.OnnxEncoder(reg.COL_ENCODER_ID)
150+
assert encoder.model_sha256 == _sha(_FAKE_MODEL)
151+
assert encoder._session is None # nothing was loaded
152+
153+
154+
def test_onnx_encoder_reports_primary_pin_from_sha256(monkeypatch):
155+
cfg = dataclasses.replace(
156+
reg.get_config(reg.COL_ENCODER_ID),
157+
sha256=_sha(_FAKE_MODEL),
158+
file_sha256=(("tokenizer.json", _sha(_FAKE_TOKENIZER)),),
159+
)
160+
monkeypatch.setitem(reg.REGISTRY, reg.COL_ENCODER_ID, cfg)
161+
assert runtime.OnnxEncoder(reg.COL_ENCODER_ID).model_sha256 == _sha(_FAKE_MODEL)
162+
163+
164+
def test_onnx_encoder_unpinned_reports_unverified(monkeypatch):
165+
cfg = dataclasses.replace(reg.get_config(reg.COL_ENCODER_ID), sha256=None, file_sha256=())
166+
monkeypatch.setitem(reg.REGISTRY, reg.COL_ENCODER_ID, cfg)
167+
assert runtime.OnnxEncoder(reg.COL_ENCODER_ID).model_sha256 == "unverified"
168+
169+
170+
def test_onnx_encoder_shipped_registry_reports_unverified():
171+
# Today's registry pins nothing, so behaviour is unchanged.
172+
for model_id in reg.REGISTRY:
173+
assert runtime.OnnxEncoder(model_id).model_sha256 == "unverified"
174+
175+
126176
def test_onnx_encoder_requires_extra():
127177
pytest.importorskip("onnxruntime")
128178
# With the extra installed but no model files, availability names the pull path.

0 commit comments

Comments
 (0)