Spaces:
Sleeping
Sleeping
File size: 772 Bytes
81f3119 aad7310 81f3119 2680899 aad7310 2680899 aad7310 2680899 81f3119 aad7310 81f3119 aad7310 81f3119 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | from sentence_transformers import SentenceTransformer
import os
_embedding_model = None
def get_embedding_model(model_name=None):
global _embedding_model
if model_name is None:
model_name = os.getenv("EMBEDDING_MODEL", "all-MiniLM-L6-v2")
if _embedding_model is None:
_embedding_model = SentenceTransformer(model_name)
return _embedding_model
def embed_chunks(chunks, model=None):
if model is None:
model = get_embedding_model()
batch_size = int(os.getenv("EMBED_BATCH_SIZE", "8"))
texts = [chunk["clean_text"] for chunk in chunks]
embeddings = model.encode(
texts,
batch_size=batch_size,
show_progress_bar=False,
convert_to_numpy=True
)
return embeddings, model
|