klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Bugfix] fix google ai embedding function (#1195)

Browse Source 
Co-authored-by: Deven Patel <deven298@yahoo.com>
This commit is contained in:
Deven Patel
2024-01-19 21:35:44 +05:30
committed by GitHub
parent d79d30bf0c
commit 0b5b12575a
1 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
embedchain/embedder/google.py
View File

@@ -1,4 +1,4 @@
from typing import Optional from typing import Optional, Union
import google.generativeai as genai import google.generativeai as genai
from chromadb import EmbeddingFunction, Embeddings from chromadb import EmbeddingFunction, Embeddings
@@ -13,12 +13,19 @@ class GoogleAIEmbeddingFunction(EmbeddingFunction):
super().__init__() super().__init__()
self.config = config or GoogleAIEmbedderConfig() self.config = config or GoogleAIEmbedderConfig()
def __call__(self, input_: str) -> Embeddings: def __call__(self, input: Union[list[str], str]) -> Embeddings:
model = self.config.model model = self.config.model
title = self.config.title title = self.config.title
task_type = self.config.task_type task_type = self.config.task_type
embeddings = genai.embed_content(model=model, content=input_, task_type=task_type, title=title) if isinstance(input, str):
return embeddings["embedding"] input_ = [input]
else:
input_ = input
data = genai.embed_content(model=model, content=input_, task_type=task_type, title=title)
embeddings = data["embedding"]
if isinstance(input_, str):
embeddings = [embeddings]
return embeddings
class GoogleAIEmbedder(BaseEmbedder): class GoogleAIEmbedder(BaseEmbedder):