from sentence_transformers import SentenceTransformer
from chromadb import Documents, Embeddings, EmbeddingFunction

class CustomEmbeddingFunction(EmbeddingFunction):
	def __call__(self, text_chunks: Documents) -> Embeddings:
		embedding_model = SentenceTransformer(
			model_name_or_path="all-mpnet-base-v2",
			device="cpu",
		)

		return embedding_model.encode(text_chunks)