Support cohere embeddings

deadbits / vigil-llm

⚡ Vigil ⚡ Detect prompt injections, jailbreaks, and other potentially risky Large Language Model (LLM) inputs

https://vigil.deadbits.ai/

Apache License 2.0

304 stars 35 forks source link

Support cohere embeddings #8

Open deadbits opened 1 year ago

deadbits commented 1 year ago

Add support for Cohere embeddings

Create datasets for Huggingface
Update vectordb scanner to use Cohere as custom embedding function
Update config parser to load Cohere model choice and API key

import cohere
co = cohere.Client(API_KEY)
response = co.embed(
  model='embed-english-v2.0',
  texts=[""])
print('Embeddings: {}'.format(response.embeddings))

ishaan-jaff commented 1 year ago

hi @deadbits I'm the maintainer of LiteLLM. We've added support for Cohere embedding models: https://docs.litellm.ai/docs/embedding/supported_embedding#cohere-embedding-models

You can call it using the OpenAI embedding Input/Output format