MANIT_Chat / server /classes /EmbeddingManager.py
WizardCoder2007's picture
first commit
2e9afea
Raw
History Blame Contribute Delete
900 Bytes
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