| """Nearest-neighbor annotation lookup against a FAISS index + parquet table |
| mounted from the HF bucket (via hfmount) -- reference_min.index / |
| reference_min.parquet, row i of the parquet matching vector i of the index. |
| Table columns: id, source (swissprot / acr / phrog), annotation, category.""" |
| import os |
|
|
| from utils.vector_store import VectorStore |
|
|
| REFERENCE_DIR = os.environ.get("REFERENCE_DATA_DIR", "/data/reference") |
| INDEX_FILENAME = os.environ.get("REFERENCE_INDEX_FILE", "reference_min.index") |
| PARQUET_FILENAME = os.environ.get("REFERENCE_PARQUET_FILE", "reference_min.parquet") |
|
|
|
|
| class AnnotationRetriever: |
| def __init__(self, data_dir=REFERENCE_DIR, index_filename=INDEX_FILENAME, |
| parquet_filename=PARQUET_FILENAME, top_k=16): |
| self.store = VectorStore( |
| os.path.join(data_dir, index_filename), |
| os.path.join(data_dir, parquet_filename), |
| id_column=None, top_k=top_k, |
| ) |
|
|
| def query(self, esm_embeddings): |
| """esm_embeddings: [N, D] array-like. Returns a list of N dicts, one |
| per query: {"top_hit", "other_source_hit"}.""" |
| scores, indices = self.store.search(esm_embeddings) |
| rows = self.store.fetch_rows(indices.ravel()) |
| return [self._summarize(s, i, rows) for s, i in zip(scores, indices)] |
|
|
| def _as_dict(self, score, row): |
| return { |
| "id": row["id"], "source": row["source"], "annotation": row["annotation"], |
| "category": row["category"], "score": float(score), |
| } |
|
|
| def _summarize(self, scores, indices, rows): |
| valid = [(float(s), int(i)) for s, i in zip(scores, indices) if i >= 0] |
| if not valid: |
| return {"top_hit": None, "other_source_hit": None} |
|
|
| top_score, top_idx = valid[0] |
| top_hit = self._as_dict(top_score, rows[top_idx]) |
| top_source = rows[top_idx]["source"] |
|
|
| other_source_hit = None |
| for score, idx in valid[1:]: |
| if rows[idx]["source"] != top_source: |
| other_source_hit = self._as_dict(score, rows[idx]) |
| break |
|
|
| return {"top_hit": top_hit, "other_source_hit": other_source_hit} |
|
|
