embed_helper.py 441 B

1234567891011
  1. #load enviroment variable
  2. from config.site import SiteConfig
  3. from sentence_transformers import SentenceTransformer
  4. config = SiteConfig()
  5. class EmbedHelper:
  6. def __init__(self):
  7. self.embedding_model_name = config.get_config("EMBEDDING_MODEL")
  8. self.embedding_model = SentenceTransformer(model_name_or_path=self.embedding_model_name)
  9. def embed_text(self, text):
  10. return self.embedding_model.encode(text).tolist()