Add OpenAI-compatible embedding API support for evaluation

README.md

@@ -178,6 +178,61 @@ alpha_ndcg, coverage, recall = evaluator.evaluate(
     results=retrieval_results,
 )
 ```
+
+### 4. Evaluate any OpenAI-compatible embedding API (e.g., Voyage) using the OpenAI SDK.
+> Make sure you install the OpenAI SDK: `pip install openai`, and export your provider key, e.g. `export VOYAGE_API_KEY=<your_key>`.
+
+```python
+import os
+
+import numpy as np
+from openai import OpenAI
+from beir.retrieval.evaluation import EvaluateRetrieval as BEIREval
+from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES
+from freshstack.retrieval.evaluation import EvaluateRetrieval
+
+# Works with Voyage, OpenAI, OpenRouter, Mistral, etc. Asymmetric models
+# (Voyage/Cohere) send an optional `input_type` via extra_body.
+class OpenAIEmbeddingModel:
+    def __init__(self, model, base_url, api_key_env, query_input_type=None, document_input_type=None):
+        self.model = model
+        self.query_input_type = query_input_type
+        self.document_input_type = document_input_type
+        self.client = OpenAI(base_url=base_url, api_key=os.environ[api_key_env])
+
+    def _embed(self, texts, input_type):
+        extra_body = {"input_type": input_type} if input_type else {}
+        resp = self.client.embeddings.create(model=self.model, input=texts, extra_body=extra_body)
+        return np.asarray([d.embedding for d in sorted(resp.data, key=lambda x: x.index)], dtype=np.float32)
+
+    def encode_queries(self, queries, batch_size=128, **kwargs):
+        return self._embed(queries, self.query_input_type)
+
+    def encode_corpus(self, corpus, batch_size=128, **kwargs):
+        texts = [(d.get("title", "") + " " + d.get("text", "")).strip() for d in corpus]
+        return self._embed(texts, self.document_input_type)
+
+model = DRES(OpenAIEmbeddingModel(
+    model="voyage-4-large",
+    base_url="https://api.voyageai.com/v1",
+    api_key_env="VOYAGE_API_KEY",
+    query_input_type="query",
+    document_input_type="document",
+), batch_size=128)
+
+retriever = BEIREval(model, score_function="cos_sim")
+retrieval_results = retriever.retrieve(corpus=corpus, queries=queries)
+
+evaluator = EvaluateRetrieval(k_values=[10, 20, 50])
+alpha_ndcg, coverage, recall = evaluator.evaluate(
+    qrels_nuggets=qrels_nuggets,
+    query_to_nuggets=query_to_nuggets,
+    qrels_query=qrels_query,
+    results=retrieval_results,
+)
+```
+> See [`examples/evaluation/api_retrieval_evaluation.py`](examples/evaluation/api_retrieval_evaluation.py) for a runnable CLI (batching + retries) that loops over all five topics.
+
 ---
 
 ## 📚 Raw Freshstack Datasets (Oct 2024)

examples/evaluation/api_retrieval_evaluation.py

@@ -0,0 +1,178 @@
+"""
+This script shows an example on how to evaluate an OpenAI-compatible embedding API (e.g., voyage-4-large) on FreshStack.
+Make sure you have the BEIR and OpenAI repositories installed: `pip install beir openai`.
+
+export VOYAGE_API_KEY=<your_key>
+for topic in langchain yolo godot angular laravel; do
+    python -m api_retrieval_evaluation \
+        --model_name_or_path "voyage-4-large" \
+        --base_url "https://api.voyageai.com/v1" \
+        --api_key_env "VOYAGE_API_KEY" \
+        --query_input_type "query" \
+        --document_input_type "document" \
+        --batch_size 128 \
+        --queries "freshstack/queries-oct-2024" \
+        --corpus "freshstack/corpus-oct-2024" \
+        --topic $topic \
+        --output_dir "./results/voyage_4_large_results/"
+done
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+import pathlib
+import time
+
+import numpy as np
+from beir import LoggingHandler
+from beir.retrieval.evaluation import EvaluateRetrieval as BEIREval
+from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES
+from openai import OpenAI
+
+from freshstack import util
+from freshstack.datasets import DataLoader
+from freshstack.retrieval.evaluation import EvaluateRetrieval
+
+#### Just some code to print debug information to stdout
+logging.basicConfig(
+    format="%(asctime)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=logging.INFO,
+    handlers=[LoggingHandler()],
+)
+#### /print debug information to stdout
+
+
+class OpenAIEmbeddingModel:
+    """OpenAI-compatible embedding API encoder for BEIR dense retrieval."""
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        base_url: str | None = None,
+        api_key_env: str = "OPENAI_API_KEY",
+        query_input_type: str | None = None,
+        document_input_type: str | None = None,
+        output_dimension: int | None = None,
+        request_batch_size: int = 128,
+        max_retries: int = 5,
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.query_input_type = query_input_type
+        self.document_input_type = document_input_type
+        self.output_dimension = output_dimension
+        self.request_batch_size = request_batch_size
+        self.max_retries = max_retries
+        self.client = OpenAI(base_url=base_url, api_key=os.environ[api_key_env])
+
+    def _embed(self, texts: list[str], input_type: str | None) -> np.ndarray:
+        extra_body = {}
+        if input_type is not None:
+            extra_body["input_type"] = input_type
+        if self.output_dimension is not None:
+            extra_body["output_dimension"] = self.output_dimension
+
+        embeddings: list[list[float]] = []
+        for start in range(0, len(texts), self.request_batch_size):
+            batch = texts[start : start + self.request_batch_size]
+            for attempt in range(self.max_retries):
+                try:
+                    response = self.client.embeddings.create(
+                        model=self.model_name_or_path, input=batch, extra_body=extra_body
+                    )
+                    ordered = sorted(response.data, key=lambda item: item.index)
+                    embeddings.extend(row.embedding for row in ordered)
+                    break
+                except Exception as e:
+                    if attempt == self.max_retries - 1:
+                        raise
+                    wait = 2**attempt
+                    logging.warning(f"Embedding request failed ({e}); retrying in {wait}s...")
+                    time.sleep(wait)
+        return np.asarray(embeddings, dtype=np.float32)
+
+    def encode_queries(self, queries: list[str], batch_size: int = 128, **kwargs) -> np.ndarray:
+        return self._embed(queries, self.query_input_type)
+
+    def encode_corpus(self, corpus: list[dict[str, str]], batch_size: int = 128, **kwargs) -> np.ndarray:
+        texts = [(doc.get("title", "") + " " + doc.get("text", "")).strip() for doc in corpus]
+        return self._embed(texts, self.document_input_type)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Compute the retrieval metrics")
+    parser.add_argument("--model_name_or_path", type=str, required=True, default="voyage-4-large")
+    parser.add_argument("--base_url", type=str, default=None, help="Base URL of the OpenAI-compatible embedding API")
+    parser.add_argument("--api_key_env", type=str, default="OPENAI_API_KEY", help="Env variable storing the API key")
+    parser.add_argument(
+        "--query_input_type", type=str, default=None, help="Input type for queries (e.g., 'query' for Voyage)"
+    )
+    parser.add_argument(
+        "--document_input_type", type=str, default=None, help="Input type for documents (e.g., 'document' for Voyage)"
+    )
+    parser.add_argument("--output_dimension", type=int, default=None, help="Optional output embedding dimension")
+    parser.add_argument(
+        "--batch_size", type=int, default=64, help="Batch size for encoding the queries and corpus documents"
+    )
+    parser.add_argument("--score_function", type=str, default="cos_sim")
+    parser.add_argument("--queries", type=str, required=True, help="The dataset to evaluate")
+    parser.add_argument(
+        "--corpus", type=str, default=None, help="The corpus dataset to evaluate, if different from the main dataset"
+    )
+    parser.add_argument("--topic", type=str, default="langchain")
+    parser.add_argument(
+        "--k_values", type=int, nargs="+", default=[5, 10, 20, 50], help="List of k values for evaluation metrics"
+    )
+    parser.add_argument("--output_dir", type=str, default=None)
+
+    args = parser.parse_args()
+
+    ### Load the nugget qrels
+    dataloader = DataLoader(queries_repo=args.queries, corpus_repo=args.corpus, topic=args.topic)
+    corpus, queries, nuggets = dataloader.load(split="test")
+    qrels_nuggets, qrels_query, query_to_nuggets = dataloader.load_qrels(split="test")
+
+    # You can also load the freshstack answers as follows:
+    # answers = dataloader.load_answers(split="test")
+
+    ### Load the OpenAI-compatible embedding model
+    model = DRES(
+        OpenAIEmbeddingModel(
+            model_name_or_path=args.model_name_or_path,
+            base_url=args.base_url,
+            api_key_env=args.api_key_env,
+            query_input_type=args.query_input_type,
+            document_input_type=args.document_input_type,
+            output_dimension=args.output_dimension,
+            request_batch_size=args.batch_size,
+        ),
+        batch_size=args.batch_size,
+    )
+
+    retriever = BEIREval(model, score_function=args.score_function)
+    results = retriever.retrieve(corpus=corpus, queries=queries)
+
+    ### Evaluate the retrieval results
+    evaluator = EvaluateRetrieval(k_values=args.k_values)
+    alpha_ndcg, coverage, recall = evaluator.evaluate(
+        qrels_nuggets=qrels_nuggets, query_to_nuggets=query_to_nuggets, qrels_query=qrels_query, results=results
+    )
+
+    ### Store the evaluation results and metrics
+    results = {query_id: results[query_id] for query_id in results if query_id in qrels_query}
+    if args.output_dir:
+        results_dir = os.path.dirname(args.output_dir)
+    else:
+        results_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "results")
+    os.makedirs(results_dir, exist_ok=True)
+
+    ### Save the evaluation results & metrics
+    util.save_runfile(os.path.join(results_dir, f"{args.topic}.run.trec"), results)
+    util.save_results(os.path.join(results_dir, f"{args.topic}.json"), alpha_ndcg, coverage, recall)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -46,6 +46,7 @@ dependencies = [
 eval = [
     "beir",
     "pylate",
+    "openai",
 ]
 dev = [
     "beir",