llm-applications

1
import os
2

3
from langchain.embeddings import OpenAIEmbeddings
4
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
5

6

7
def get_embedding_model(embedding_model_name, model_kwargs, encode_kwargs):
8
    if embedding_model_name == "text-embedding-ada-002":
9
        embedding_model = OpenAIEmbeddings(
10
            model=embedding_model_name,
11
            openai_api_base=os.environ["OPENAI_API_BASE"],
12
            openai_api_key=os.environ["OPENAI_API_KEY"],
13
        )
14
    else:
15
        embedding_model = HuggingFaceEmbeddings(
16
            model_name=embedding_model_name,
17
            model_kwargs=model_kwargs,
18
            encode_kwargs=encode_kwargs,
19
        )
20
    return embedding_model
21

22

23
class EmbedChunks:
24
    def __init__(self, model_name):
25
        # Embedding model
26
        self.embedding_model = get_embedding_model(
27
            embedding_model_name=model_name,
28
            model_kwargs={"device": "cuda"},
29
            encode_kwargs={"device": "cuda", "batch_size": 100},
30
        )
31

32
    def __call__(self, batch):
33
        embeddings = self.embedding_model.embed_documents(batch["text"])
34
        return {"text": batch["text"], "source": batch["source"], "embeddings": embeddings}
35