FormuLLaMa-Demo / utils /embedding_utils.py
MarioCerulo's picture
Upload 3 files
50b2d56 verified
raw
history blame contribute delete
376 Bytes
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)