Return Cohere embeddings as lists of floats (#1394)

This PR fixes the types returned by Cohere embeddings. Currently, Cohere
client returns instances of `cohere.embeddings.Embeddings`. Since the
transport layer relies on JSON, some numbers might be represented as
ints, not floats, which happens quite often. While that doesn't seem to
be an issue, it breaks some pydantic models if they require strict
floats.
searx-doc
Kacper Łukawski 1 year ago committed by GitHub
parent 1989e7d4c2
commit 8947797250
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -64,7 +64,7 @@ class CohereEmbeddings(BaseModel, Embeddings):
embeddings = self.client.embed( embeddings = self.client.embed(
model=self.model, texts=texts, truncate=self.truncate model=self.model, texts=texts, truncate=self.truncate
).embeddings ).embeddings
return embeddings return [list(map(float, e)) for e in embeddings]
def embed_query(self, text: str) -> List[float]: def embed_query(self, text: str) -> List[float]:
"""Call out to Cohere's embedding endpoint. """Call out to Cohere's embedding endpoint.
@ -78,4 +78,4 @@ class CohereEmbeddings(BaseModel, Embeddings):
embedding = self.client.embed( embedding = self.client.embed(
model=self.model, texts=[text], truncate=self.truncate model=self.model, texts=[text], truncate=self.truncate
).embeddings[0] ).embeddings[0]
return embedding return list(map(float, embedding))

Loading…
Cancel
Save