vector_index.py

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