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

add cohere encoder

parent ce56320a
Branches
Tags
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.
Please register or to comment