From 1f626345f9e63e4ee4416b5ef61f80a128a58525 Mon Sep 17 00:00:00 2001
From: "Ziyang \"Claude\" Hu" <33562602+ClaudeHu@users.noreply.github.com>
Date: Thu, 9 Nov 2023 15:09:02 -0500
Subject: [PATCH] refactor sentencetransformers to fastembedding

---
 bbconf/bbconf.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/bbconf/bbconf.py b/bbconf/bbconf.py
index 984df0b..d1f9795 100644
--- a/bbconf/bbconf.py
+++ b/bbconf/bbconf.py
@@ -52,7 +52,7 @@
 # os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"  # to suppress verbose warnings tensorflow
 from geniml.text2bednn import text2bednn
 from geniml.search import QdrantBackend
-from sentence_transformers import SentenceTransformer
+from fastembed.embedding import FlagEmbedding
 from geniml.region2vec import Region2VecExModel
 from geniml.io import RegionSet
 
@@ -510,8 +510,8 @@ def _create_t2bsi_object(self) -> Union[text2bednn.Text2BEDSearchInterface, None
 
         try:
             return text2bednn.Text2BEDSearchInterface(
-                nl2vec_model=SentenceTransformer(
-                    os.getenv("HF_MODEL", DEFAULT_HF_MODEL)
+                nl2vec_model=FlagEmbedding(
+                    model_name=os.getenv("HF_MODEL", DEFAULT_HF_MODEL)
                 ),
                 vec2vec_model=self._config[CFG_PATH_KEY][CFG_PATH_VEC2VEC_KEY],
                 search_backend=self.qdrant_backend,
@@ -700,3 +700,4 @@ def get_drs_metadata(
         )
 
         return drs_dict
+