From f2abdd09c871c5220a83da84898581e8133efd2d Mon Sep 17 00:00:00 2001
From: "Ziyang \"Claude\" Hu" <33562602+ClaudeHu@users.noreply.github.com>
Date: Thu, 9 Nov 2023 15:05:34 -0500
Subject: [PATCH] refactor sentence transformer to fastembedding

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

diff --git a/bbconf/bbconf.py b/bbconf/bbconf.py
index 984df0b..a6f7c55 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,