my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	from typing import Optional, Union
#3
#4	from chromadb import EmbeddingFunction, Embeddings
#5
#6	from embedchain.config import BaseEmbedderConfig
#7	from embedchain.embedder.base import BaseEmbedder
#8
#9
#10	class ClarifaiEmbeddingFunction(EmbeddingFunction):
#11	def __init__(self, config: BaseEmbedderConfig) -> None:
#12	super().__init__()
#13	try:
#14	from clarifai.client.input import Inputs
#15	from clarifai.client.model import Model
#16	except ModuleNotFoundError:
#17	raise ModuleNotFoundError(
#18	"The required dependencies for ClarifaiEmbeddingFunction are not installed."
#19	'Please install with `pip install --upgrade "embedchain[clarifai]"`'
#20	) from None
#21	self.config = config
#22	self.api_key = config.api_key or os.getenv("CLARIFAI_PAT")
#23	self.model = config.model
#24	self.model_obj = Model(url=self.model, pat=self.api_key)
#25	self.input_obj = Inputs(pat=self.api_key)
#26
#27	def __call__(self, input: Union[str, list[str]]) -> Embeddings:
#28	if isinstance(input, str):
#29	input = [input]
#30
#31	batch_size = 32
#32	embeddings = []
#33	try:
#34	for i in range(0, len(input), batch_size):
#35	batch = input[i : i + batch_size]
#36	input_batch = [
#37	self.input_obj.get_text_input(input_id=str(id), raw_text=inp) for id, inp in enumerate(batch)
#38	]
#39	response = self.model_obj.predict(input_batch)
#40	embeddings.extend([list(output.data.embeddings[0].vector) for output in response.outputs])
#41	except Exception as e:
#42	print(f"Predict failed, exception: {e}")
#43
#44	return embeddings
#45
#46
#47	class ClarifaiEmbedder(BaseEmbedder):
#48	def __init__(self, config: Optional[BaseEmbedderConfig] = None):
#49	super().__init__(config)
#50
#51	embedding_func = ClarifaiEmbeddingFunction(config=self.config)
#52	self.set_embedding_fn(embedding_fn=embedding_func)
#53

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public