add support for google gemini embeddings as an embedder (#56)

* wip on retrieve

* add ir measures

* add support for google embeddings

* voyage api key

* update embeddings and update reqs

* remove todo

* github head ref

* fix workflow

* revert workflow

---------

Co-authored-by: Mihail Eric <mihaileric@Mihails-MacBook-Pro.local>

benchmarks/retrieval/requirements.txt

@@ -0,0 +1,2 @@
+dotenv
+ir_measures

benchmarks/retrieval/retrieve.py

@@ -2,22 +2,24 @@
 
 Make sure to `pip install ir_measures` before running this script.
 """
-
 import json
 import logging
 import os
 import time
 
 import configargparse
+from dotenv import load_dotenv
 from ir_measures import MAP, MRR, P, Qrel, R, Rprec, ScoredDoc, calc_aggregate, nDCG
 
 import sage.config
+from sage.data_manager import GitHubRepoManager
 from sage.retriever import build_retriever_from_args
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 
+load_dotenv()
 
 def main():
     parser = configargparse.ArgParser(
@@ -35,17 +37,27 @@ def main():
         default=None,
         help="Path where to output predictions and metrics. Optional, since metrics are also printed to console.",
     )
-    parser.add("--max-instances", default=None, type=int, help="Maximum number of instances to process.")
 
+    parser.add("--max-instances", default=None, type=int, help="Maximum number of instances to process.")
     sage.config.add_config_args(parser)
     sage.config.add_llm_args(parser)  # Needed for --multi-query-retriever, which rewrites the query with an LLM.
     sage.config.add_embedding_args(parser)
     sage.config.add_vector_store_args(parser)
     sage.config.add_reranking_args(parser)
+    sage.config.add_repo_args(parser)
+    sage.config.add_indexing_args(parser)
     args = parser.parse_args()
     sage.config.validate_vector_store_args(args)
-
-    retriever = build_retriever_from_args(args)
+    repo_manager = GitHubRepoManager(
+            args.repo_id,
+            commit_hash=args.commit_hash,
+            access_token=os.getenv("GITHUB_TOKEN"),
+            local_dir=args.local_dir,
+            inclusion_file=args.include,
+            exclusion_file=args.exclude,
+    )
+    repo_manager.download()
+    retriever = build_retriever_from_args(args, repo_manager)
 
     with open(args.benchmark, "r") as f:
         benchmark = json.load(f)
@@ -70,7 +82,7 @@ def main():
         item["retrieved"] = []
         for doc_idx, doc in enumerate(retrieved):
             # The absolute value of the scores below does not affect the metrics; it merely determines the ranking of
-            # the retrived documents. The key of the score varies depending on the underlying retriever. If there's no
+            # the retrieved documents. The key of the score varies depending on the underlying retriever. If there's no
             # score, we use 1/(doc_idx+1) since it preserves the order of the documents.
             score = doc.metadata.get("score", doc.metadata.get("relevance_score", 1 / (doc_idx + 1)))
             retrieved_docs.append(ScoredDoc(query_id=query_id, doc_id=doc.metadata["file_path"], score=score))
@@ -83,7 +95,6 @@ def main():
     print("Calculating metrics...")
     results = calc_aggregate([Rprec, P @ 1, R @ 3, nDCG @ 3, MAP, MRR], golden_docs, retrieved_docs)
     results = {str(key): value for key, value in results.items()}
-
     if args.logs_dir:
         if not os.path.exists(args.logs_dir):
             os.makedirs(args.logs_dir)

requirements.txt

@@ -3,6 +3,7 @@ Pygments==2.18.0
 cohere==5.9.2
 configargparse
 fastapi==0.112.2
+google-ai-generativelanguage==0.6.6
 gradio>=4.26.0
 langchain==0.2.16
 langchain-anthropic==0.1.23
@@ -10,6 +11,7 @@ langchain-cohere==0.2.4
 langchain-community==0.2.17
 langchain-core==0.2.41
 langchain-experimental==0.0.65
+langchain-google-genai
 langchain-nvidia-ai-endpoints==0.2.2
 langchain-ollama==0.1.3
 langchain-openai==0.1.25

sage/config.py

@@ -11,6 +11,11 @@ from configargparse import ArgumentParser
 
 from sage.reranker import RerankerProvider
 
+# Limits defined here: https://ai.google.dev/gemini-api/docs/models/gemini
+# NOTE: MAX_CHUNKS_PER_BATCH isn't documented anywhere but we pick a reasonable value
+GEMINI_MAX_CHUNKS_PER_BATCH = 64
+GEMINI_MAX_TOKENS_PER_CHUNK = 2048
+
 MARQO_MAX_CHUNKS_PER_BATCH = 64
 # The ADA embedder from OpenAI has a maximum of 8192 tokens.
 OPENAI_MAX_TOKENS_PER_CHUNK = 8192
@@ -82,7 +87,7 @@ def add_repo_args(parser: ArgumentParser) -> Callable:
 
 def add_embedding_args(parser: ArgumentParser) -> Callable:
     """Adds embedding-related arguments to the parser and returns a validator."""
-    parser.add("--embedding-provider", default="marqo", choices=["openai", "voyage", "marqo"])
+    parser.add("--embedding-provider", default="marqo", choices=["openai", "voyage", "marqo", "gemini"])
     parser.add(
         "--embedding-model",
         type=str,
@@ -304,6 +309,26 @@ def _validate_marqo_embedding_args(args):
         )
 
 
+def _validate_gemini_embedding_args(args):
+    """Validates the configuration of the Gemini batch embedder and sets defaults."""
+    if not args.embedding_model:
+        args.embedding_model = "models/text-embedding-004"
+    assert os.environ["GOOGLE_API_KEY"], "Please set the GOOGLE_API_KEY environment variable if using `gemini` embeddings."
+    if not args.chunks_per_batch:
+        args.chunks_per_batch = GEMINI_MAX_CHUNKS_PER_BATCH
+    elif args.chunks_per_batch > GEMINI_MAX_CHUNKS_PER_BATCH:
+        args.chunks_per_batch = GEMINI_MAX_CHUNKS_PER_BATCH
+        logging.warning(
+            f"Gemini enforces a limit of {GEMINI_MAX_CHUNKS_PER_BATCH} chunks per batch. "
+            "Overwriting embeddings.chunks_per_batch."
+        )
+
+    if not args.tokens_per_chunk:
+        args.tokens_per_chunk = GEMINI_MAX_TOKENS_PER_CHUNK
+    if not args.embedding_size:
+        args.embedding_size = 768
+
+
 def validate_embedding_args(args):
     """Validates the configuration of the batch embedder and sets defaults."""
     if args.embedding_provider == "openai":
@@ -312,6 +337,8 @@ def validate_embedding_args(args):
         _validate_voyage_embedding_args(args)
     elif args.embedding_provider == "marqo":
         _validate_marqo_embedding_args(args)
+    elif args.embedding_provider == "gemini":
+        _validate_gemini_embedding_args(args)
     else:
         raise ValueError(f"Unrecognized --embedding-provider={args.embedding_provider}")
 

sage/embedder.py

@@ -4,16 +4,25 @@ import json
 import logging
 import os
 import time
-from abc import ABC, abstractmethod
+from abc import ABC
+from abc import abstractmethod
 from collections import Counter
-from typing import Dict, Generator, List, Optional, Tuple
+from typing import Dict
+from typing import Generator
+from typing import List
+from typing import Optional
+from typing import Tuple
 
+import google.generativeai as genai
 import marqo
 import requests
 from openai import OpenAI
-from tenacity import retry, stop_after_attempt, wait_random_exponential
+from tenacity import retry
+from tenacity import stop_after_attempt
+from tenacity import wait_random_exponential
 
-from sage.chunker import Chunk, Chunker
+from sage.chunker import Chunk
+from sage.chunker import Chunker
 from sage.constants import TEXT_FIELD
 from sage.data_manager import DataManager
 
@@ -63,7 +72,7 @@ class OpenAIBatchEmbedder(BatchEmbedder):
 
             if len(batch) > chunks_per_batch:
                 for i in range(0, len(batch), chunks_per_batch):
-                    sub_batch = batch[i : i + chunks_per_batch]
+                    sub_batch = batch[i: i + chunks_per_batch]
                     openai_batch_id = self._issue_job_for_chunks(sub_batch, batch_id=f"{dataset_name}/{len(batch_ids)}")
                     batch_ids[openai_batch_id] = [chunk.metadata for chunk in sub_batch]
                     if max_embedding_jobs and len(batch_ids) >= max_embedding_jobs:
@@ -233,7 +242,7 @@ class VoyageBatchEmbedder(BatchEmbedder):
 
             if len(batch) > chunks_per_batch:
                 for i in range(0, len(batch), chunks_per_batch):
-                    sub_batch = batch[i : i + chunks_per_batch]
+                    sub_batch = batch[i: i + chunks_per_batch]
                     logging.info("Embedding %d chunks...", len(sub_batch))
                     result = self._make_batch_request(sub_batch)
                     for chunk, datum in zip(sub_batch, result["data"]):
@@ -305,7 +314,7 @@ class MarqoEmbedder(BatchEmbedder):
 
             if len(batch) > chunks_per_batch:
                 for i in range(0, len(batch), chunks_per_batch):
-                    sub_batch = batch[i : i + chunks_per_batch]
+                    sub_batch = batch[i: i + chunks_per_batch]
                     logging.info("Indexing %d chunks...", len(sub_batch))
                     self.index.add_documents(
                         documents=[chunk.metadata for chunk in sub_batch],
@@ -335,6 +344,79 @@ class MarqoEmbedder(BatchEmbedder):
         return []
 
 
+class GeminiBatchEmbedder(BatchEmbedder):
+    """Batch embedder that calls Gemini."""
+
+    def __init__(self, data_manager: DataManager, chunker: Chunker, embedding_model: str):
+        self.data_manager = data_manager
+        self.chunker = chunker
+        self.embedding_data = []
+        self.embedding_model = embedding_model
+        genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+
+    def _make_batch_request(self, chunks: List[Chunk]) -> Dict:
+        return genai.embed_content(
+            model=self.embedding_model,
+            content=[chunk.content for chunk in chunks],
+            task_type="retrieval_document")
+
+    def embed_dataset(self, chunks_per_batch: int, max_embedding_jobs: int = None):
+        """Issues batch embedding jobs for the entire dataset."""
+        batch = []
+        chunk_count = 0
+
+        request_count = 0
+        last_request_time = time.time()
+
+        for content, metadata in self.data_manager.walk():
+            chunks = self.chunker.chunk(content, metadata)
+            chunk_count += len(chunks)
+            batch.extend(chunks)
+
+            if len(batch) > chunks_per_batch:
+                for i in range(0, len(batch), chunks_per_batch):
+                    sub_batch = batch[i: i + chunks_per_batch]
+                    logging.info("Embedding %d chunks...", len(sub_batch))
+                    result = self._make_batch_request(sub_batch)
+                    for chunk, embedding in zip(sub_batch, result["embedding"]):
+                        self.embedding_data.append((chunk.metadata, embedding))
+                    request_count += 1
+
+                    # Check if we've made more than 1500 requests in the last minute
+                    # Rate limits here: https://ai.google.dev/gemini-api/docs/models/gemini
+                    current_time = time.time()
+                    elapsed_time = current_time - last_request_time
+                    if elapsed_time < 60 and request_count >= 1400:
+                        logging.info("Reached rate limit, pausing for 60 seconds...")
+                        time.sleep(60)
+                        last_request_time = current_time
+                        request_count = 0
+                    # Reset the last request time and request count if more than 60 sec have passed
+                    elif elapsed_time > 60:
+                        last_request_time = current_time
+                        request_count = 0
+
+                batch = []
+
+        # Finally, commit the last batch.
+        if batch:
+            logging.info("Embedding %d chunks...", len(batch))
+            result = self._make_batch_request(batch)
+            for chunk, embedding in zip(batch, result["embedding"]):
+                self.embedding_data.append((chunk.metadata, embedding))
+
+        logging.info(f"Successfully embedded {chunk_count} chunks.")
+
+    def embeddings_are_ready(self, *args, **kwargs) -> bool:
+        """Checks whether the batch embedding jobs are done."""
+        return True
+
+    def download_embeddings(self, *args, **kwargs) -> Generator[Vector, None, None]:
+        """Yields (chunk_metadata, embedding) pairs for each chunk in the dataset."""
+        for chunk_metadata, embedding in self.embedding_data:
+            yield chunk_metadata, embedding
+
+
 def build_batch_embedder_from_flags(data_manager: DataManager, chunker: Chunker, args) -> BatchEmbedder:
     if args.embedding_provider == "openai":
         return OpenAIBatchEmbedder(data_manager, chunker, args.local_dir, args.embedding_model, args.embedding_size)
@@ -344,5 +426,7 @@ def build_batch_embedder_from_flags(data_manager: DataManager, chunker: Chunker,
         return MarqoEmbedder(
             data_manager, chunker, index_name=args.index_namespace, url=args.marqo_url, model=args.embedding_model
         )
+    elif args.embedding_provider == "gemini":
+        return GeminiBatchEmbedder(data_manager, chunker, embedding_model=args.embedding_model)
     else:
         raise ValueError(f"Unrecognized embedder type {args.embedding_provider}")

sage/reranker.py

@@ -10,7 +10,6 @@ from langchain_core.documents import BaseDocumentCompressor
 from langchain_nvidia_ai_endpoints import NVIDIARerank
 from langchain_voyageai import VoyageAIRerank
 
-
 class RerankerProvider(Enum):
     NONE = "none"
     HUGGINGFACE = "huggingface"

sage/retriever.py

@@ -1,24 +1,30 @@
+from typing import Optional
+
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain.retrievers.multi_query import MultiQueryRetriever
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
 from langchain_openai import OpenAIEmbeddings
 from langchain_voyageai import VoyageAIEmbeddings
 
+from sage.data_manager import DataManager
 from sage.llm import build_llm_via_langchain
 from sage.reranker import build_reranker
 from sage.vector_store import build_vector_store_from_args
 
 
-def build_retriever_from_args(args):
+def build_retriever_from_args(args, data_manager: Optional[DataManager] = None):
     """Builds a retriever (with optional reranking) from command-line arguments."""
 
     if args.embedding_provider == "openai":
         embeddings = OpenAIEmbeddings(model=args.embedding_model)
     elif args.embedding_provider == "voyage":
         embeddings = VoyageAIEmbeddings(model=args.embedding_model)
+    elif args.embedding_provider == "gemini":
+        embeddings = GoogleGenerativeAIEmbeddings(model=args.embedding_model)
     else:
         embeddings = None
 
-    retriever = build_vector_store_from_args(args).as_retriever(
+    retriever = build_vector_store_from_args(args, data_manager).as_retriever(
         top_k=args.retriever_top_k, embeddings=embeddings, namespace=args.index_namespace
     )
 

sage/vector_store.py

@@ -75,7 +75,6 @@ class PineconeVectorStore(VectorStore):
         self.dimension = dimension
         self.client = Pinecone()
         self.alpha = alpha
-
         if alpha < 1.0:
             if bm25_cache and os.path.exists(bm25_cache):
                 logging.info("Loading BM25 encoder from cache.")
@@ -192,9 +191,14 @@ def build_vector_store_from_args(args: dict, data_manager: Optional[DataManager]
     """
     if args.vector_store_provider == "pinecone":
         bm25_cache = os.path.join(".bm25_cache", args.index_namespace, "bm25_encoder.json")
-
-        if not os.path.exists(bm25_cache) and data_manager:
+        if args.retrieval_alpha < 1.0 and not os.path.exists(bm25_cache) and data_manager:
             logging.info("Fitting BM25 encoder on the corpus...")
+            if is_punkt_downloaded():
+                print("punkt is already downloaded")
+            else:
+                print("punkt is not downloaded")
+                # Optionally download it
+                nltk.download('punkt_tab')
             corpus = [content for content, _ in data_manager.walk()]
             bm25_encoder = BM25Encoder()
             bm25_encoder.fit(corpus)