92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
|
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",))
|