Skip to content
Snippets Groups Projects
Unverified Commit 48e2330e authored by James Briggs's avatar James Briggs
Browse files

add cohere encoder

parent ce56320a
No related branches found
No related tags found
No related merge requests found
import os
import cohere
from decision_layer.encoders import BaseEncoder
class CohereEncoder(BaseEncoder):
def __init__(self, name: str):
super().__init__(name)
client: cohere.Client | None
def __init__(self, name: str, cohere_api_key: str | None = None):
super().__init__(name=name, client=None)
cohere_api_key = cohere_api_key or os.getenv("COHERE_API_KEY")
if cohere_api_key is None:
raise ValueError("Cohere API key cannot be 'None'.")
self.client = cohere.Client(cohere_api_key)
def __call__(self, texts: list[str]) -> list[float]:
raise NotImplementedError
\ No newline at end of file
if len(texts) == 1:
input_type = "search_query"
else:
input_type = "search_document"
embeds = self.client.embed(
texts, input_type=input_type, model="embed-english-v3.0"
)
return embeds.embeddings
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment