pluto
74e12eca60
FastAPI service that loads ViT-L-14 model and exposes /embed, /similarity, and /batch-similarity endpoints. Runs on RTX 3090 with NVIDIA GPU reservation. Includes 9 cosine similarity tests. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
131 lines
4.8 KiB
Python
131 lines
4.8 KiB
Python
"""Tests for CLIP service math utilities.
|
|
|
|
These tests exercise the cosine similarity logic using pure numpy --
|
|
no GPU, no model loading, no network calls required.
|
|
"""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
# Import the cosine_similarity function directly from the app module.
|
|
# We add the app directory to sys.path so the import works without
|
|
# installing the package.
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(
|
|
0, os.path.join(os.path.dirname(__file__), "..", "app")
|
|
)
|
|
|
|
from main import cosine_similarity
|
|
|
|
|
|
class TestCosineSimilarity:
|
|
"""Test the cosine_similarity helper function."""
|
|
|
|
def test_identical_vectors(self):
|
|
"""Identical vectors should have cosine similarity of 1.0."""
|
|
v = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
|
|
score = cosine_similarity(v, v)
|
|
assert score == pytest.approx(1.0, abs=1e-6)
|
|
|
|
def test_identical_unit_vectors(self):
|
|
"""Identical unit vectors should also yield 1.0."""
|
|
v = np.array([0.0, 1.0, 0.0], dtype=np.float32)
|
|
score = cosine_similarity(v, v)
|
|
assert score == pytest.approx(1.0, abs=1e-6)
|
|
|
|
def test_orthogonal_vectors(self):
|
|
"""Orthogonal vectors should have cosine similarity of 0.0."""
|
|
a = np.array([1.0, 0.0, 0.0], dtype=np.float32)
|
|
b = np.array([0.0, 1.0, 0.0], dtype=np.float32)
|
|
score = cosine_similarity(a, b)
|
|
assert score == pytest.approx(0.0, abs=1e-6)
|
|
|
|
def test_opposite_vectors(self):
|
|
"""Opposite vectors should have cosine similarity of -1.0."""
|
|
a = np.array([1.0, 2.0, 3.0], dtype=np.float32)
|
|
b = np.array([-1.0, -2.0, -3.0], dtype=np.float32)
|
|
score = cosine_similarity(a, b)
|
|
assert score == pytest.approx(-1.0, abs=1e-6)
|
|
|
|
def test_zero_vector_returns_zero(self):
|
|
"""A zero vector should return 0.0 (guarded against division by zero)."""
|
|
a = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
|
b = np.array([1.0, 2.0, 3.0], dtype=np.float32)
|
|
score = cosine_similarity(a, b)
|
|
assert score == pytest.approx(0.0, abs=1e-6)
|
|
|
|
def test_high_dimensional(self):
|
|
"""Test with 768-dimensional vectors (CLIP ViT-L/14 output size)."""
|
|
rng = np.random.RandomState(42)
|
|
a = rng.randn(768).astype(np.float32)
|
|
# Similarity with itself should be 1.0
|
|
score = cosine_similarity(a, a)
|
|
assert score == pytest.approx(1.0, abs=1e-5)
|
|
|
|
|
|
class TestBatchSimilarity:
|
|
"""Test batch similarity logic (matching the endpoint's numpy math)."""
|
|
|
|
def _batch_similarity(
|
|
self, query: np.ndarray, candidates: np.ndarray
|
|
) -> tuple[list[float], int, float]:
|
|
"""Replicate the batch-similarity endpoint logic."""
|
|
query_norm = query / (np.linalg.norm(query) + 1e-10)
|
|
candidate_norms = candidates / (
|
|
np.linalg.norm(candidates, axis=1, keepdims=True) + 1e-10
|
|
)
|
|
scores = candidate_norms @ query_norm
|
|
scores_list = [float(s) for s in scores]
|
|
best_index = int(np.argmax(scores))
|
|
best_score = scores_list[best_index]
|
|
return scores_list, best_index, best_score
|
|
|
|
def test_best_match_index(self):
|
|
"""The candidate most similar to the query should be identified."""
|
|
query = np.array([1.0, 0.0, 0.0], dtype=np.float32)
|
|
candidates = np.array(
|
|
[
|
|
[0.0, 1.0, 0.0], # orthogonal
|
|
[0.9, 0.1, 0.0], # close match
|
|
[1.0, 0.0, 0.0], # exact match
|
|
[0.5, 0.5, 0.0], # partial match
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
scores, best_idx, best_score = self._batch_similarity(query, candidates)
|
|
|
|
assert best_idx == 2
|
|
assert best_score == pytest.approx(1.0, abs=1e-5)
|
|
assert len(scores) == 4
|
|
# Orthogonal candidate should score ~0
|
|
assert scores[0] == pytest.approx(0.0, abs=1e-5)
|
|
|
|
def test_single_candidate(self):
|
|
"""Batch with a single candidate should work correctly."""
|
|
query = np.array([1.0, 1.0], dtype=np.float32)
|
|
candidates = np.array([[1.0, 1.0]], dtype=np.float32)
|
|
|
|
scores, best_idx, best_score = self._batch_similarity(query, candidates)
|
|
assert best_idx == 0
|
|
assert best_score == pytest.approx(1.0, abs=1e-5)
|
|
|
|
def test_scores_ordering(self):
|
|
"""Scores should reflect relative similarity to the query."""
|
|
query = np.array([1.0, 0.0], dtype=np.float32)
|
|
candidates = np.array(
|
|
[
|
|
[0.5, 0.866], # ~60 degrees
|
|
[0.866, 0.5], # ~30 degrees
|
|
[1.0, 0.0], # 0 degrees (exact)
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
scores, best_idx, _ = self._batch_similarity(query, candidates)
|
|
|
|
# Exact match should be best
|
|
assert best_idx == 2
|
|
# 30 degrees should score higher than 60 degrees
|
|
assert scores[1] > scores[0]
|