refactor sentence transformer to fastembedding

databio · Nov 9, 2023 · f2abdd0 · f2abdd0
1 parent 974b698
commit f2abdd0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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,