Update retriever.py

retriever.py

@@ -1,39 +1,67 @@
 from huggingface_hub import hf_hub_download
+from FlagEmbedding import FlagICLModel
 import pandas as pd
 import faiss
 import numpy as np
 
 class FoundationRetriever:
-    def __init__(self, repo_id, parquet_file, faiss_file, bge_model="BAAI/bge-m3", use_fp16=False):
-        # 1️⃣ Download files from HF hub
+    """
+    Retrieve foundations whose missions align with a user perspective.
+    Uses FAISS for efficient similarity search and BGE-en-ICL for encoding queries with examples.
+    """
+
+    def __init__(self, repo_id, parquet_file, faiss_file, use_fp16=False):
+        """
+        repo_id: Hugging Face dataset repo containing Parquet and FAISS index
+        parquet_file: Parquet filename in repo
+        faiss_file: FAISS index filename in repo
+        use_fp16: whether to use FP16 for model inference
+        """
+        # Download files from Hugging Face
         parquet_path = hf_hub_download(repo_id, parquet_file)
         faiss_path = hf_hub_download(repo_id, faiss_file)
 
-        # 2️⃣ Load parquet
+        # Load foundations dataframe
         self.df = pd.read_parquet(parquet_path)
 
-        # 3️⃣ Load FAISS index
+        # Load FAISS index
         self.index = faiss.read_index(faiss_path)
         self.dim = self.index.d
 
-        # 4️⃣ Load BGE-M3 model
-        from FlagEmbedding import FlagICLModel
+        # Define few-shot examples for ICL query embedding
+        examples_for_task = [
+            {
+                "instruct": "Retrieve foundations whose mission aligns with the given perspective.",
+                "query": "Protect marine life while educating children about ocean conservation",
+                "response": "Foundations working on marine conservation and youth education."
+            },
+            {
+                "instruct": "Retrieve foundations whose mission aligns with the given perspective.",
+                "query": "Promote renewable energy education and community awareness",
+                "response": "Foundations focused on clean energy advocacy and public education."
+            }
+        ]
+
+        # Load BGE-en-ICL model
         self.model = FlagICLModel(
-            bge_model,
-            query_instruction_for_retrieval="Given a user perspective, retrieve relevant foundation missions.",
-            examples_for_task=None,
+            "BAAI/bge-en-icl",
+            query_instruction_for_retrieval="Given a user perspective about philanthropy, retrieve relevant foundation missions.",
+            examples_for_task=examples_for_task,
             use_fp16=use_fp16
         )
 
     def find_aligned_foundations(self, perspective_text, top_k=5):
-        # Encode perspective
+        """
+        Returns a DataFrame with top-k foundations aligned with the user's perspective.
+        """
+        # Encode user perspective with ICL
         query_emb = self.model.encode_queries([perspective_text]).astype("float32")
         faiss.normalize_L2(query_emb)
 
         # Search FAISS
         distances, indices = self.index.search(query_emb, top_k)
 
-        # Retrieve foundations
+        # Retrieve foundation info
         results = self.df.iloc[indices[0]][["Title", "Purpose"]].copy()
         results["similarity"] = distances[0]