from sentence_transformers import SentenceTransformer
from langchain.schema import Document



class EmbeddingGen:

    def __init__(self, model_name: str):
        self.model = SentenceTransformer(model_name)
    
    def embed_documents(self, chunks):
        return [self.model.encode(chunk) for chunk in chunks]
    
    def embed_query(self, text):
        return self.model.encode(text)