category retrieval

utils/annotation_retrieval.py

@@ -1,7 +1,7 @@
 """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 has a 'source' column (swissprot / acr / phrogs)."""
+Table columns: id, source (swissprot / acr / phrog), annotation, category."""
 import os
 
 import faiss
@@ -42,7 +42,10 @@ class AnnotationRetriever:
 
     def _as_dict(self, score, idx):
         row = self.df.iloc[idx]
-        return {"id": row["id"], "source": row["source"], "annotation": row["annotation"], "score": float(score)}
+        return {
+            "id": row["id"], "source": row["source"], "annotation": row["annotation"],
+            "category": row["category"], "score": float(score),
+        }
 
     def _summarize(self, scores, indices):
         valid = [(float(s), int(i)) for s, i in zip(scores, indices) if i >= 0]