Spaces:
Sleeping
Sleeping
File size: 3,030 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 | import os
import pandas as pd
import chromadb
from sentence_transformers import SentenceTransformer
_COLLECTION_NAME = "asl_gloss"
_MODEL_NAME = "all-MiniLM-L6-v2"
_instance: "VectorIndex | None" = None
class VectorIndex:
def __init__(
self,
csv_path: str = "content/asl_app_data/asl_video_index_final_with_path_cleaned.csv",
chroma_dir: str | None = None,
):
if chroma_dir is None:
chroma_dir = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db")
self._model = SentenceTransformer(_MODEL_NAME)
client = chromadb.PersistentClient(path=chroma_dir)
self._collection = client.get_or_create_collection(
_COLLECTION_NAME,
metadata={"hnsw:space": "cosine"},
)
if self._collection.count() == 0:
self._build(csv_path)
else:
print(f"VectorIndex: reusing existing collection ({self._collection.count()} entries)")
def _build(self, csv_path: str) -> None:
df = pd.read_csv(csv_path)
seen: dict[str, str] = {}
for _, row in df.iterrows():
fname = os.path.basename(str(row["path"]))
for col in ("token", "phrase", "word"):
val = row.get(col, "")
if pd.notna(val) and str(val).strip():
key = str(val).upper().strip()
if key not in seen:
seen[key] = fname
keys = list(seen.keys())
fnames = list(seen.values())
print(f"VectorIndex: encoding {len(keys)} tokens…")
embeddings = self._model.encode(keys, show_progress_bar=False).tolist()
self._collection.upsert(
ids=keys,
embeddings=embeddings,
metadatas=[{"filename": f} for f in fnames],
documents=keys,
)
print(f"VectorIndex: indexed {len(keys)} tokens")
def query_gloss(self, token: str, top_k: int = 3) -> list[tuple[str, float, str]]:
threshold = float(os.getenv("SEMANTIC_THRESHOLD", "0.6"))
embedding = self._model.encode([token]).tolist()
results = self._collection.query(
query_embeddings=embedding,
n_results=top_k,
include=["metadatas", "distances", "documents"],
)
out: list[tuple[str, float, str]] = []
for key, distance, meta in zip(
results["documents"][0],
results["distances"][0],
results["metadatas"][0],
):
score = 1.0 - distance # ChromaDB cosine space: distance = 1 - similarity
if score >= threshold:
out.append((key, score, meta["filename"]))
return out
def get_vector_index(
csv_path: str = "content/asl_app_data/asl_video_index_final_with_path_cleaned.csv",
) -> "VectorIndex":
"""Return the process-level singleton VectorIndex, creating it on first call."""
global _instance
if _instance is None:
_instance = VectorIndex(csv_path=csv_path)
return _instance
|