:gem: [Feature] Embedder: check_model_name and switch_model, and separate constants

configs/constants.py

@@ -0,0 +1,3 @@
+AVAILABLE_MODELS = [
+    "jinaai/jina-embeddings-v2-base-zh",
+]

transforms/embed.py

@@ -7,6 +7,7 @@ from transformers import AutoModel
 from numpy.linalg import norm
 
 from configs.envs import ENVS
+from configs.constants import AVAILABLE_MODELS
 
 os.environ["HF_ENDPOINT"] = ENVS["HF_ENDPOINT"]
 os.environ["HF_TOKEN"] = ENVS["HF_TOKEN"]
@@ -17,13 +18,24 @@ def cosine_similarity(a, b):
 
 
 class JinaAIEmbedder:
-    def __init__(self, model_name: str = "jinaai/jina-embeddings-v2-base-zh"):
+    def __init__(self, model_name: str = AVAILABLE_MODELS[0]):
         self.model_name = model_name
         self.load_model()
 
+    def check_model_name(self):
+        if self.model_name not in AVAILABLE_MODELS:
+            self.model_name = AVAILABLE_MODELS[0]
+        return True
+
     def load_model(self):
+        self.check_model_name()
         self.model = AutoModel.from_pretrained(self.model_name, trust_remote_code=True)
 
+    def switch_model(self, model_name: str):
+        if model_name != self.model_name:
+            self.model_name = model_name
+            self.load_model()
+
     def encode(self, text: Union[str, list[str]]):
         if isinstance(text, str):
             text = [text]