File size: 900 Bytes
2e9afea
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from sentence_transformers import SentenceTransformer
from typing import List,Dict,Any,Tuple
import numpy as np

class EmbeddingManager:
    def __init__(self,model_name: str= "BAAI/bge-large-en-v1.5"):
        self.model_name= model_name
        self.model= None
        self._load_model()

    def _load_model(self):
        try:
            print(f"Embedding model: {self.model_name}")
            self.model= SentenceTransformer(self.model_name)
            print(f"suceess in loading model, embedding dimensions: {self.model.get_sentence_embedding_dimension()}")
        except Exception as e:
            print("error in loading model")
            raise
    
    def generate_embeddings(self,texts: List[str])-> np.ndarray:
        if not self.model:
            raise ValueError("model not found")
        embeddings= self.model.encode(texts,show_progress_bar= True)
        return embeddings