Spaces:
Sleeping
Sleeping
File size: 3,679 Bytes
dd590f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import os
import numpy as np
import pytest
from unittest.mock import MagicMock, patch
@pytest.fixture
def csv_file(tmp_path):
p = tmp_path / "index.csv"
p.write_text(
"path,token,phrase,word\n"
"video_dir/hello.mp4,HELLO,,\n"
"video_dir/thank_you.mp4,THANK-YOU,,\n"
)
return str(p)
@pytest.fixture
def chroma_dir(tmp_path):
return str(tmp_path / "chroma")
def _empty_client():
col = MagicMock()
col.count.return_value = 0
client = MagicMock()
client.get_or_create_collection.return_value = col
return client, col
def _full_client():
col = MagicMock()
col.count.return_value = 50
client = MagicMock()
client.get_or_create_collection.return_value = col
return client, col
def test_builds_index_when_collection_empty(csv_file, chroma_dir):
"""On first init, upserts all CSV token keys into ChromaDB."""
client, col = _empty_client()
model = MagicMock()
model.encode.return_value = np.array([[0.1, 0.2], [0.3, 0.4]])
with patch("chromadb.PersistentClient", return_value=client), \
patch("vector_index.SentenceTransformer", return_value=model):
from vector_index import VectorIndex
VectorIndex(csv_path=csv_file, chroma_dir=chroma_dir)
col.upsert.assert_called_once()
ids = col.upsert.call_args.kwargs["ids"]
assert "HELLO" in ids
assert "THANK-YOU" in ids
def test_skips_build_when_collection_populated(csv_file, chroma_dir):
"""Does not re-encode if collection already has data."""
client, col = _full_client()
model = MagicMock()
with patch("chromadb.PersistentClient", return_value=client), \
patch("vector_index.SentenceTransformer", return_value=model):
from vector_index import VectorIndex
VectorIndex(csv_path=csv_file, chroma_dir=chroma_dir)
col.upsert.assert_not_called()
def test_query_gloss_returns_above_threshold(csv_file, chroma_dir):
"""Returns (key, score, filename) tuples for results above SEMANTIC_THRESHOLD."""
client, col = _full_client()
col.query.return_value = {
"documents": [["HELP"]],
"distances": [[0.15]], # score = 1 - 0.15 = 0.85
"metadatas": [[{"filename": "help.mp4"}]],
}
model = MagicMock()
model.encode.return_value = np.array([[0.1, 0.2]])
with patch("chromadb.PersistentClient", return_value=client), \
patch("vector_index.SentenceTransformer", return_value=model):
from vector_index import VectorIndex
idx = VectorIndex(csv_path=csv_file, chroma_dir=chroma_dir)
with patch.dict(os.environ, {"SEMANTIC_THRESHOLD": "0.6"}):
results = idx.query_gloss("ASSIST")
assert len(results) == 1
key, score, fname = results[0]
assert key == "HELP"
assert abs(score - 0.85) < 0.001
assert fname == "help.mp4"
def test_query_gloss_filters_below_threshold(csv_file, chroma_dir):
"""Returns [] when best cosine score is below SEMANTIC_THRESHOLD."""
client, col = _full_client()
col.query.return_value = {
"documents": [["HELP"]],
"distances": [[0.7]], # score = 1 - 0.7 = 0.3 < 0.6
"metadatas": [[{"filename": "help.mp4"}]],
}
model = MagicMock()
model.encode.return_value = np.array([[0.1, 0.2]])
with patch("chromadb.PersistentClient", return_value=client), \
patch("vector_index.SentenceTransformer", return_value=model):
from vector_index import VectorIndex
idx = VectorIndex(csv_path=csv_file, chroma_dir=chroma_dir)
with patch.dict(os.environ, {"SEMANTIC_THRESHOLD": "0.6"}):
results = idx.query_gloss("XYZZY")
assert results == []
|