File size: 1,086 Bytes
6158da4
f0018f2
1e2550f
6158da4
 
 
 
 
 
 
f51bb92
6158da4
 
f51bb92
6158da4
 
57b7b8d
6158da4
 
 
f51bb92
6d056d5
f51bb92
6d056d5
 
 
6158da4
b83cc65
6158da4
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
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
from modules.config.constants import OPENAI_API_KEY, HUGGINGFACE_TOKEN


class EmbeddingModelLoader:
    def __init__(self, config):
        self.config = config

    def load_embedding_model(self):
        if self.config["vectorstore"]["model"] in ["text-embedding-ada-002"]:
            embedding_model = OpenAIEmbeddings(
                deployment="SL-document_embedder",
                model=self.config["vectorestore"]["model"],
                show_progress_bar=True,
                openai_api_key=OPENAI_API_KEY,
                disallowed_special=(),
            )
        else:
            embedding_model = HuggingFaceEmbeddings(
                model_name=self.config["vectorstore"]["model"],
                model_kwargs={
                    "device": f"{self.config['device']}",
                    "token": f"{HUGGINGFACE_TOKEN}",
                    "trust_remote_code": True,
                },
            )

        return embedding_model