pinecone.data.features.inference.inference_asyncio

  1from typing import Optional, Dict, List, Union, Any
  2
  3from pinecone.core.openapi.inference.api.inference_api import AsyncioInferenceApi
  4from .models import EmbeddingsList, RerankResult
  5
  6from .inference_request_builder import (
  7    InferenceRequestBuilder,
  8    EmbedModel as EmbedModelEnum,
  9    RerankModel as RerankModelEnum,
 10)
 11
 12
 13class AsyncioInference:
 14    """
 15    The `AsyncioInference` class configures and uses the Pinecone Inference API to generate embeddings and
 16    rank documents.
 17
 18    This class is generally not instantiated directly, but rather accessed through a parent `Pinecone` client
 19    object that is responsible for managing shared configurations.
 20
 21    ```python
 22    from pinecone import PineconeAsyncio
 23
 24    pc = PineconeAsyncio()
 25    embeddings = await pc.inference.embed(
 26        model="text-embedding-3-small",
 27        inputs=["Hello, world!"],
 28        parameters={"input_type": "passage", "truncate": "END"}
 29    )
 30    ```
 31
 32    :param config: A `pinecone.config.Config` object, configured and built in the Pinecone class.
 33    :type config: `pinecone.config.Config`, required
 34    """
 35
 36    EmbedModel = EmbedModelEnum
 37    RerankModel = RerankModelEnum
 38
 39    def __init__(self, api_client, **kwargs) -> None:
 40        self.api_client = api_client
 41        """ @private """
 42
 43        self.__inference_api = AsyncioInferenceApi(api_client)
 44        """ @private """
 45
 46    async def embed(
 47        self,
 48        model: str,
 49        inputs: Union[str, List[Dict], List[str]],
 50        parameters: Optional[Dict[str, Any]] = None,
 51    ) -> EmbeddingsList:
 52        """
 53        Generates embeddings for the provided inputs using the specified model and (optional) parameters.
 54
 55        :param model: The model to use for generating embeddings.
 56        :type model: str, required
 57
 58        :param inputs: A list of items to generate embeddings for.
 59        :type inputs: list, required
 60
 61        :param parameters: A dictionary of parameters to use when generating embeddings.
 62        :type parameters: dict, optional
 63
 64        :return: EmbeddingsList object with keys `data`, `model`, and `usage`. The `data` key contains a list of
 65        `n` embeddings, where `n` = len(inputs) and type(n) = Embedding. Precision of returned embeddings is either
 66        float16 or float32, with float32 being the default. `model` key is the model used to generate the embeddings.
 67        `usage` key contains the total number of tokens used at request-time.
 68
 69        Example:
 70        >>> inputs = ["Who created the first computer?"]
 71        >>> outputs = await pc.inference.embed(model="multilingual-e5-large", inputs=inputs, parameters={"input_type": "passage", "truncate": "END"})
 72        >>> print(outputs)
 73        EmbeddingsList(
 74            model='multilingual-e5-large',
 75            data=[
 76                {'values': [0.1, ...., 0.2]},
 77              ],
 78            usage={'total_tokens': 6}
 79        )
 80        """
 81        request_body = InferenceRequestBuilder.embed_request(
 82            model=model, inputs=inputs, parameters=parameters
 83        )
 84        resp = await self.__inference_api.embed(embed_request=request_body)
 85        return EmbeddingsList(resp)
 86
 87    async def rerank(
 88        self,
 89        model: str,
 90        query: str,
 91        documents: Union[List[str], List[Dict[str, Any]]],
 92        rank_fields: List[str] = ["text"],
 93        return_documents: bool = True,
 94        top_n: Optional[int] = None,
 95        parameters: Optional[Dict[str, Any]] = None,
 96    ) -> RerankResult:
 97        """
 98        Rerank documents with associated relevance scores that represent the relevance of each document
 99        to the provided query using the specified model.
100
101        :param model: The model to use for reranking.
102        :type model: str, required
103
104        :param query: The query to compare with documents.
105        :type query: str, required
106
107        :param documents: A list of documents or strings to rank.
108        :type documents: list, required
109
110        :param rank_fields: A list of document fields to use for ranking. Defaults to ["text"].
111        :type rank_fields: list, optional
112
113        :param return_documents: Whether to include the documents in the response. Defaults to True.
114        :type return_documents: bool, optional
115
116        :param top_n: How many documents to return. Defaults to len(documents).
117        :type top_n: int, optional
118
119        :param parameters: A dictionary of parameters to use when ranking documents.
120        :type parameters: dict, optional
121
122        :return: RerankResult object with keys `data` and `usage`. The `data` key contains a list of
123        `n` documents, where `n` = `top_n` and type(n) = Document. The documents are sorted in order of
124        relevance, with the first being the most relevant. The `index` field can be used to locate the document
125        relative to the list of documents specified in the request. Each document contains a `score` key
126        representing how close the document relates to the query.
127
128        Example:
129        >>> result = await pc.inference.rerank(
130                model="bge-reranker-v2-m3",
131                query="Tell me about tech companies",
132                documents=[
133                    "Apple is a popular fruit known for its sweetness and crisp texture.",
134                    "Software is still eating the world.",
135                    "Many people enjoy eating apples as a healthy snack.",
136                    "Acme Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces.",
137                    "An apple a day keeps the doctor away, as the saying goes.",
138                ],
139                top_n=2,
140                return_documents=True,
141            )
142        >>> print(result)
143        RerankResult(
144          model='bge-reranker-v2-m3',
145          data=[
146            { index=3, score=0.020980744,
147              document={text="Acme Inc. has rev..."} },
148            { index=1, score=0.00034015716,
149              document={text="Software is still..."} }
150          ],
151          usage={'rerank_units': 1}
152        )
153        """
154        rerank_request = InferenceRequestBuilder.rerank(
155            model=model,
156            query=query,
157            documents=documents,
158            rank_fields=rank_fields,
159            return_documents=return_documents,
160            top_n=top_n,
161            parameters=parameters,
162        )
163        resp = await self.__inference_api.rerank(rerank_request=rerank_request)
164        return RerankResult(resp)