update retriever to add inference option

retriever.py

@@ -1,35 +1,64 @@
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, InferenceClient
 from FlagEmbedding import FlagICLModel
 import pandas as pd
 import faiss
 import numpy as np
+import os
+import json
 
-class FoundationRetriever:
-    """
-    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)
-
-        # Load foundations dataframe
-        self.df = pd.read_parquet(parquet_path)
-
-        # Load FAISS index
-        self.index = faiss.read_index(faiss_path)
-        self.dim = self.index.d
-
-        # Define few-shot examples for ICL query embedding
-        examples_for_task = [
+# Define HF tokens
+HF_TOKEN_read = os.environ.get("HF_TOKEN_read")
+HF_TOKEN_inference = os.environ.get("HF_TOKEN_inf")
+
+# Setup InferenceClient for embeddings
+client = InferenceClient(provider="nebius", api_key=HF_TOKEN_inference)
+
+# Dataset repo (private)
+DATASET_REPO = "luciagomez/MrPhil_vector"
+
+# -------------------------------------------------------------------
+# 1. Download files from Hugging Face dataset
+# -------------------------------------------------------------------
+parquet_path = hf_hub_download(
+    repo_id=DATASET_REPO,
+    filename="foundations.parquet",
+    repo_type="dataset",
+    token=HF_TOKEN_read
+)
+
+faiss_path = hf_hub_download(
+    repo_id=DATASET_REPO,
+    filename="faiss.index",
+    repo_type="dataset",
+    token=HF_TOKEN_read
+)
+
+meta_path = hf_hub_download(
+    repo_id=DATASET_REPO,
+    filename="meta.json",
+    repo_type="dataset",
+    token=HF_TOKEN_read
+)
+
+
+# -------------------------------------------------------------------
+# 2. Load data
+# -------------------------------------------------------------------
+df = pd.read_parquet(parquet_path)
+index = faiss.read_index(faiss_path)
+with open(meta_path, "r") as f:
+    meta = json.load(f)
+
+dim = meta.get("embedding_dim", None)
+n = meta.get("num_vectors", len(df))
+
+print(f"Loaded FAISS index with {n} vectors of dimension {dim}")
+
+
+# -------------------------------------------------------------------
+# 3. Initialize BGE-ICL model for queries
+# -------------------------------------------------------------------
+examples = [
             {
                 "instruct": "Retrieve foundations whose mission aligns with the given perspective.",
                 "query": "Protect marine life while educating children about ocean conservation",
@@ -42,27 +71,32 @@ class FoundationRetriever:
             }
         ]
 
-        # Load BGE-en-ICL model
-        self.model = FlagICLModel(
-            "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):
-        """
-        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 foundation info
-        results = self.df.iloc[indices[0]][["Title", "Purpose"]].copy()
-        results["similarity"] = distances[0]
-
-        return results
+
+#model = FlagICLModel(
+#    "BAAI/bge-en-icl",
+#    query_instruction_for_retrieval="Given a mission statement, retrieve foundations with aligned purposes.",
+#    examples_for_task=examples,
+#    use_fp16=False
+#)
+
+def encode_query(perspective):
+    payload = {"model": "BAAI/bge-en-icl", 
+               "inputs": perspective, 
+               "parameters": {"instruction": "Retrieve foundations aligned with perspective.", 
+                              "examples": examples}}
+    response = client.feature_extraction(**payload)
+    return np.array(response)
+
+# -------------------------------------------------------------------
+# 4. Retrieval function
+# -------------------------------------------------------------------
+def find_similar_foundations(perspective, top_k=5):
+    q_emb = encode_query(perspective).astype("float32")
+    faiss.normalize_L2(q_emb)
+    scores, idxs = index.search(q_emb, top_k)
+    return [
+        {"title": df.iloc[i]["Title"], "purpose": df.iloc[i]["Purpose"], "similarity": float(scores[0][j])}
+        for j, i in enumerate(idxs[0])
+    ]
+
+