File size: 649 Bytes
e04dd70
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import numpy as np
from langchain.embeddings.openai import OpenAIEmbeddings
from utils import ModelName


class EmbeddingManager:
    def __init__(self, model_name=ModelName.OPENAI) -> None:
        self.model_name = model_name

    def compare_embeddigns_similarity(self, embedding_1, embedding_2):
        similarity = np.dot(embedding_1, embedding_2)
        return similarity

    def generate_embeddings(self, splits: list[str]):
        embedding = None
        if self.model_name == ModelName.OPENAI:
            embedding = OpenAIEmbeddings()
        embeddings = [embedding.embed_query(split) for split in splits]
        return embeddings