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