Skip to content

Add support for OpenAI, Azure, Cohere, Replicate Embedding Models - using litellm #15

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: dev
Choose a base branch
from

Conversation

ishaan-jaff
Copy link

@ishaan-jaff ishaan-jaff commented Aug 6, 2023

This PR adds support for models from all the above mentioned providers using https://github.com/BerriAI/litellm/
TLDR

  • You get support for all liteLLM supported embedding models
  • I noticed you had a class for Cohere embeddings, instead of maintaining classes, liteLLM can be the simple abstraction to select your models. Simply pass the model param when calling embedding and liteLLM manages calling the correct API & translation of input + output

Here's a sample of how it's used:

from litellm import embedding

## set ENV variables
# ENV variables can be set in .env file, too. Example in .env.example
os.environ["OPENAI_API_KEY"] = "openai key"
os.environ["COHERE_API_KEY"] = "cohere key"

# openai call
response = embedding(input=[text], model="text-embedding-ada-002")

# azure call
response = embedding(input=[text], model="azure-embedding-model", azure=True)

# replicate call
response = embedding(input=[text], model="replicate/all-mpnet-base-v2:b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305")


@ishaan-jaff
Copy link
Author

Cc @mishanti @FelixKirsch can you take a look at this when you get a chance ? Happy to make additonal changes/update docs based on your feedback

You can also get logging for your embedding calls made through liteLLM😊

input=documents_batch, engine=self.model_name
)
else:
response = openai.Embedding.create(
response = embedding(
input=documents_batch, model=self.model_name
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in order to use cohere/replicate/other embedding models, just pass in the the model here

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant