Cohere embeddings generation missing async batch feature. How to speed up generation?

[burberg92]

Hi, I was looking to work with cohere embeddings and I really need to speed up the embedding generation. After looking at llamaindex I noticed that it does not support async batch generation.

these are the current functions for generating embeddings for cohere:

def _get_query_embedding(self, query: str) -> List[float]:
    """Get query embedding."""
    return self._embed([query])[0]

async def _aget_query_embedding(self, query: str) -> List[float]:
    """Get query embedding async."""
    return self._get_query_embedding(query)

def _get_text_embedding(self, text: str) -> List[float]:
    """Get text embedding."""
    return self._embed([text])[0]

async def _aget_text_embedding(self, text: str) -> List[float]:
    """Get text embedding async."""
    return self._get_text_embedding(text)

def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
    """Get text embeddings."""
    return self._embed(texts)

I believe that it is missing an async version of get_text_embeddings. This would greatly speed up the generation of embeddings. How can I implement this? or what strategies would be recommended to speed up the generation of embeddings for cohere?

Alternatively, I see that in the BaseEmbedding file there is a function called aget_text_embedding_batch, which is what I want to do. How can I use this?

I only want to generate the embeddings from a list of text, not create an index. Thank you.

[dosubot[bot]]

🤖

Hello @burberg92, nice to see you again! I hope you're doing well.

Based on the information you've provided, it seems you're looking to implement an async version of get_text_embeddings to speed up the generation of Cohere embeddings. You're correct that the current implementation of LlamaIndex does not support async batch generation. However, you can leverage the aget_text_embedding_batch function from the BaseEmbedding file to achieve this.

Here's how you can use it:

class YourEmbedding(BaseEmbedding):
    # Your implementation details...

    async def aget_text_embedding_batch(
        self, texts: List[str], show_progress: bool = False
    ) -> List[Embedding]:
        return await super().aget_text_embedding_batch(texts, show_progress)

In this code, YourEmbedding is your subclass of BaseEmbedding. The aget_text_embedding_batch method is overridden to call the same method from the superclass (BaseEmbedding). This will use the asynchronous batch text embedding method from the base class.

Remember, you need to implement the _aget_text_embedding method in your subclass, as it's an abstract method in the base class and is used by aget_text_embedding_batch.

class YourEmbedding(BaseEmbedding):
    # Your implementation details...

    async def _aget_text_embedding(self, text: str) -> Embedding:
        # Your async implementation here...

This way, you can speed up the generation of embeddings by processing multiple texts concurrently. You can find more details about this in the BaseEmbedding file.

Regarding your question about the _embed method, it's not explicitly defined in the provided context for the OpenAIEmbedding class in the LlamaIndex codebase. However, there are several methods related to embedding, such as _get_query_embedding, _aget_query_embedding, _get_text_embedding, _aget_text_embedding, _get_text_embeddings, and _aget_text_embeddings. If you are referring to a different _embed method, please provide more context or check the correct file or class in the LlamaIndex codebase.

I hope this helps! If you have any other questions or need further clarification, feel free to ask.

Sources

• Why do we send embedding requests to Open AI one by one (one for each chunk) instead of batching them in a few requests?
• llama_index/embeddings/base.py
• llama_index/embeddings/openai.py

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.