This commit is contained in:
1 parent
4db3f45164
commit
d25c04cdd2
175 files changed
+56243
-15
No files matched your search
@@ -0,0 +1,92 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from gewu.adapters.embeddings import (
|
||||
EmbeddingProviderError,
|
||||
LocalSentenceTransformer,
|
||||
model_artifact_sha256,
|
||||
)
|
||||
|
||||
|
||||
class FakeEncoder:
|
||||
def __init__(self, rows: list[list[float]]) -> None:
|
||||
self.rows = rows
|
||||
self.seen: list[str] = []
|
||||
self.calls = 0
|
||||
|
||||
def encode(
|
||||
self,
|
||||
sentences: list[str],
|
||||
*,
|
||||
normalize_embeddings: bool,
|
||||
convert_to_numpy: bool,
|
||||
show_progress_bar: bool,
|
||||
) -> Any:
|
||||
assert normalize_embeddings and convert_to_numpy and not show_progress_bar
|
||||
self.seen.extend(sentences)
|
||||
result = self.rows[self.calls : self.calls + len(sentences)]
|
||||
self.calls += len(sentences)
|
||||
return result
|
||||
|
||||
|
||||
def _model_path(tmp_path: Path) -> Path:
|
||||
path = tmp_path / "model"
|
||||
path.mkdir()
|
||||
(path / "config.json").write_text('{"dim":2}', encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
def test_local_sentence_transformer_applies_e5_prefixes_and_metadata(tmp_path: Path) -> None:
|
||||
path = _model_path(tmp_path)
|
||||
encoder = FakeEncoder([[0.6, 0.8], [1.0, 0.0]])
|
||||
artifact = model_artifact_sha256(path)
|
||||
provider = LocalSentenceTransformer(
|
||||
path,
|
||||
model_id="intfloat/multilingual-e5-small@pinned",
|
||||
expected_sha256=artifact,
|
||||
dimensions=2,
|
||||
encoder=encoder,
|
||||
)
|
||||
|
||||
assert provider.embed_queries(("question",)) == ((0.6, 0.8),)
|
||||
assert provider.embed_passages(("evidence",)) == ((1.0, 0.0),)
|
||||
assert encoder.seen == ["query: question", "passage: evidence"]
|
||||
assert provider.model_id == "intfloat/multilingual-e5-small@pinned"
|
||||
assert provider.artifact_sha256 == artifact
|
||||
assert len(provider.artifact_sha256) == 64
|
||||
|
||||
|
||||
def test_local_sentence_transformer_rejects_changed_artifact(tmp_path: Path) -> None:
|
||||
path = _model_path(tmp_path)
|
||||
(path / "config.json").write_text("changed", encoding="utf-8")
|
||||
|
||||
with pytest.raises(EmbeddingProviderError, match="digest"):
|
||||
LocalSentenceTransformer(
|
||||
path,
|
||||
model_id="model",
|
||||
expected_sha256=hashlib.sha256(b"wrong").hexdigest(),
|
||||
dimensions=2,
|
||||
encoder=FakeEncoder([[1.0, 0.0]]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rows", [[[2.0, 0.0]], [[float("nan"), 0.0]], [[1.0]]])
|
||||
def test_local_sentence_transformer_rejects_invalid_vectors(
|
||||
tmp_path: Path, rows: list[list[float]]
|
||||
) -> None:
|
||||
path = _model_path(tmp_path)
|
||||
provider = LocalSentenceTransformer(
|
||||
path,
|
||||
model_id="model",
|
||||
expected_sha256=model_artifact_sha256(path),
|
||||
dimensions=2,
|
||||
encoder=FakeEncoder(rows),
|
||||
)
|
||||
|
||||
with pytest.raises(EmbeddingProviderError):
|
||||
provider.embed_passages(("text",))
|
||||
Reference in new issue
Block a user