Merge pull request #1062 from stanfordnlp/revert-1027-main

Revert "fix(pinecone_rm): refactored to use cloud_embed and fix pinecone init"

Author: arnavsinghvi11

dspy/retrieve/pinecone_rm.py

@@ -3,7 +3,6 @@
 Author: Dhar Rawal (@drawal1)
 """
 
-from abc import ABC, abstractmethod
 from typing import List, Optional, Union
 
 import backoff
@@ -14,6 +13,9 @@
 try:
     import pinecone
 except ImportError:
+    pinecone = None
+
+if pinecone is None:
     raise ImportError(
         "The pinecone library is required to use PineconeRM. Install it with `pip install dspy-ai[pinecone]`",
     )
@@ -31,64 +33,6 @@
 except Exception:
     ERRORS = (openai.RateLimitError, openai.APIError)
 
-
-class CloudEmbedProvider(ABC):
-    def __init__ (self, model, api_key=None):
-        self.model = model
-        self.api_key = api_key
-    
-    @abstractmethod
-    def get_embeddings(self, queries: List[str]) -> List[List[float]]:
-        pass
-
-class OpenAIEmbed(CloudEmbedProvider):
-    def __init__(self, model="text-embedding-ada-002", api_key: Optional[str]=None, org: Optional[str]=None):
-        super().__init__(model, api_key)
-        self.org = org
-        if self.api_key:
-            openai.api_key = self.api_key
-        if self.org:
-            openai.organization = org
-
-    
-    @backoff.on_exception(
-        backoff.expo,
-        ERRORS,
-        max_time=15,
-    )
-    def get_embeddings(self, queries: List[str]) -> List[List[float]]:
-        if OPENAI_LEGACY:
-            embedding = openai.Embedding.create(
-                input=queries, model=self.model,
-            )
-        else:
-            embedding = openai.embeddings.create(
-                input=queries, model=self.model,
-            ).model_dump()
-        return [embedding["embedding"] for embedding in embedding["data"]]
-    
-class CohereEmbed(CloudEmbedProvider):
-    def __init__(self, model: str = "multilingual-22-12", api_key: Optional[str] = None):
-        try:
-            import cohere
-        except ImportError:
-            raise ImportError(
-                "The cohere library is required to use CohereEmbed. Install it with `pip install cohere`",
-            )
-        super().__init__(model, api_key)
-        self.client = cohere.Client(api_key)
-    
-    @backoff.on_exception(
-        backoff.expo,
-        ERRORS,
-        max_time=15,
-    )
-    def get_embeddings(self, queries: List[str]) -> List[List[float]]:
-        embeddings = self.client.embed(texts=queries, model=self.model).embeddings
-        return embeddings
-    
-
-
 class PineconeRM(dspy.Retrieve):
     """
     A retrieval module that uses Pinecone to return the top passages for a given query.
@@ -99,8 +43,11 @@ class PineconeRM(dspy.Retrieve):
     Args:
         pinecone_index_name (str): The name of the Pinecone index to query against.
         pinecone_api_key (str, optional): The Pinecone API key. Defaults to None.
+        pinecone_env (str, optional): The Pinecone environment. Defaults to None.
         local_embed_model (str, optional): The local embedding model to use. A popular default is "sentence-transformers/all-mpnet-base-v2".
-        cloud_emded_provider (CloudEmbedProvider, optional): The cloud embedding provider to use. Defaults to None.
+        openai_embed_model (str, optional): The OpenAI embedding model to use. Defaults to "text-embedding-ada-002".
+        openai_api_key (str, optional): The API key for OpenAI. Defaults to None.
+        openai_org (str, optional): The organization for OpenAI. Defaults to None.
         k (int, optional): The number of top passages to retrieve. Defaults to 3.
 
     Returns:
@@ -110,7 +57,6 @@ class PineconeRM(dspy.Retrieve):
         Below is a code snippet that shows how to use this as the default retriver:
         ```python
         llm = dspy.OpenAI(model="gpt-3.5-turbo")
-        retriever_model = PineconeRM(index_name, cloud_emded_provider=OpenAIEmbed())
         retriever_model = PineconeRM(openai.api_key)
         dspy.settings.configure(lm=llm, rm=retriever_model)
         ```
@@ -125,8 +71,11 @@ def __init__(
         self,
         pinecone_index_name: str,
         pinecone_api_key: Optional[str] = None,
+        pinecone_env: Optional[str] = None,
         local_embed_model: Optional[str] = None,
-        cloud_emded_provider: Optional[CloudEmbedProvider] = None,
+        openai_embed_model: Optional[str] = "text-embedding-ada-002",
+        openai_api_key: Optional[str] = None,
+        openai_org: Optional[str] = None,
         k: int = 3,
     ):
         if local_embed_model is not None:
@@ -146,25 +95,69 @@ def __init__(
                 'mps' if torch.backends.mps.is_available()
                 else 'cpu',
             )
-
-        elif cloud_emded_provider is not None:
+        elif openai_embed_model is not None:
+            self._openai_embed_model = openai_embed_model
             self.use_local_model = False
-            self.cloud_emded_provider = cloud_emded_provider
-        
+            # If not provided, defaults to env vars OPENAI_API_KEY and OPENAI_ORGANIZATION
+            if openai_api_key:
+                openai.api_key = openai_api_key
+            if openai_org:
+                openai.organization = openai_org
         else:
             raise ValueError(
-                "Either local_embed_model or cloud_embed_provider must be provided.",
+                "Either local_embed_model or openai_embed_model must be provided.",
             )
 
-        if pinecone_api_key is None:
-            self.pinecone_client = pinecone.Pinecone()
-        else:
-            self.pinecone_client = pinecone.Pinecone(api_key=pinecone_api_key)
-
-        self._pinecone_index = self.pinecone_client.Index(pinecone_index_name)
+        self._pinecone_index = self._init_pinecone(
+            pinecone_index_name, pinecone_api_key, pinecone_env,
+        )
 
         super().__init__(k=k)
 
+    def _init_pinecone(
+        self,
+        index_name: str,
+        api_key: Optional[str] = None,
+        environment: Optional[str] = None,
+        dimension: Optional[int] = None,
+        distance_metric: Optional[str] = None,
+    ) -> pinecone.Index:
+        """Initialize pinecone and return the loaded index.
+
+        Args:
+            index_name (str): The name of the index to load. If the index is not does not exist, it will be created.
+            api_key (str, optional): The Pinecone API key, defaults to env var PINECONE_API_KEY if not provided.
+            environment (str, optional): The environment (ie. `us-west1-gcp` or `gcp-starter`. Defaults to env PINECONE_ENVIRONMENT.
+
+        Raises:
+            ValueError: If api_key or environment is not provided and not set as an environment variable.
+
+        Returns:
+            pinecone.Index: The loaded index.
+        """
+
+        # Pinecone init overrides default if kwargs are present, so we need to exclude if None
+        kwargs = {}
+        if api_key:
+            kwargs["api_key"] = api_key
+        if environment:
+            kwargs["environment"] = environment
+        pinecone.init(**kwargs)
+
+        active_indexes = pinecone.list_indexes()
+        if index_name not in active_indexes:
+            if dimension is None and distance_metric is None:
+                raise ValueError(
+                    "dimension and distance_metric must be provided since the index provided does not exist.",
+                )
+
+            pinecone.create_index(
+                name=index_name,
+                dimension=dimension,
+                metric=distance_metric,
+            )
+
+        return pinecone.Index(index_name)
 
     def _mean_pooling(
             self, 
@@ -182,7 +175,11 @@ def _mean_pooling(
         input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
         return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
 
-
+    @backoff.on_exception(
+        backoff.expo,
+        ERRORS,
+        max_time=15,
+    )
     def _get_embeddings(
         self, 
         queries: List[str],
@@ -195,16 +192,24 @@ def _get_embeddings(
         Returns:
             List[List[float]]: List of embeddings corresponding to each query.
         """
-        if not self.use_local_model:
-            return self.cloud_emded_provider.get_embeddings(queries)
-
         try:
             import torch
         except ImportError as exc:
             raise ModuleNotFoundError(
                 "You need to install torch to use a local embedding model with PineconeRM.",
             ) from exc
 
+        if not self.use_local_model:
+            if OPENAI_LEGACY:
+                embedding = openai.Embedding.create(
+                    input=queries, model=self._openai_embed_model,
+                )
+            else:
+                embedding = openai.embeddings.create(
+                    input=queries, model=self._openai_embed_model,
+                ).model_dump()
+            return [embedding["embedding"] for embedding in embedding["data"]]
+        
         # Use local model
         encoded_input = self._local_tokenizer(queries, padding=True, truncation=True, return_tensors="pt").to(self.device)
         with torch.no_grad():
@@ -217,55 +222,51 @@ def _get_embeddings(
         # we need a pooling strategy to get a single vector representation of the input
         # so the default is to take the mean of the hidden states
 
-    def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = None) -> dspy.Prediction:
+    def forward(self, query_or_queries: Union[str, List[str]]) -> dspy.Prediction:
         """Search with pinecone for self.k top passages for query
 
         Args:
             query_or_queries (Union[str, List[str]]): The query or queries to search for.
-            k (Optional[int]): The number of top passages to retrieve. Defaults to self.k
 
         Returns:
             dspy.Prediction: An object containing the retrieved passages.
         """
-        k = k if k is not None else self.k
         queries = (
             [query_or_queries]
             if isinstance(query_or_queries, str)
             else query_or_queries
         )
         queries = [q for q in queries if q]  # Filter empty queries
         embeddings = self._get_embeddings(queries)
+
         # For single query, just look up the top k passages
         if len(queries) == 1:
             results_dict = self._pinecone_index.query(
-                vector=embeddings[0], top_k=k, include_metadata=True,
+                embeddings[0], top_k=self.k, include_metadata=True,
             )
 
             # Sort results by score
             sorted_results = sorted(
                 results_dict["matches"], key=lambda x: x.get("scores", 0.0), reverse=True,
             )
-            
             passages = [result["metadata"]["text"] for result in sorted_results]
-            passages = [dotdict({"long_text": passage}) for passage in passages]
-            return passages
+            passages = [dotdict({"long_text": passage for passage in passages})]
+            return dspy.Prediction(passages=passages)
 
         # For multiple queries, query each and return the highest scoring passages
         # If a passage is returned multiple times, the score is accumulated. For this reason we increase top_k by 3x
         passage_scores = {}
         for embedding in embeddings:
             results_dict = self._pinecone_index.query(
-                vector=embedding, top_k=k * 3, include_metadata=True,
+                embedding, top_k=self.k * 3, include_metadata=True,
             )
             for result in results_dict["matches"]:
                 passage_scores[result["metadata"]["text"]] = (
                     passage_scores.get(result["metadata"]["text"], 0.0)
                     + result["score"]
                 )
-        
+
         sorted_passages = sorted(
             passage_scores.items(), key=lambda x: x[1], reverse=True,
-        )[: k]
-        
-        passages=[dotdict({"long_text": passage}) for passage, _ in sorted_passages]
-        return passages
+        )[: self.k]
+        return dspy.Prediction(passages=[dotdict({"long_text": passage}) for passage, _ in sorted_passages])