BUG: fix custom embedding launch error (#2016)

xorbitsai · Aug 5, 2024 · 1112993 · 1112993
1 parent ec421ed
commit 1112993
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/xinference/model/embedding/core.py b/xinference/model/embedding/core.py
@@ -151,8 +151,8 @@ def to(self, *args, **kwargs):
 
         patch_trust_remote_code()
         if (
-            "gte-Qwen2" in self._model_spec.model_id
-            or "gte-Qwen2" in self._model_spec.model_name
+            "gte" in self._model_spec.model_name.lower()
+            and "qwen2" in self._model_spec.model_name.lower()
         ):
             self._model = XSentenceTransformer(
                 self._model_path,
@@ -260,8 +260,8 @@ def encode(
                 device = model._target_device
 
             if (
-                "gte-Qwen2" not in self._model_spec.model_id
-                and "gte-Qwen2" not in self._model_spec.model_name
+                "gte" in self._model_spec.model_name.lower()
+                and "qwen2" in self._model_spec.model_name.lower()
             ):
                 model.to(device)
 
@@ -342,8 +342,8 @@ def encode(
             return all_embeddings, all_token_nums
 
         if (
-            "gte-Qwen2" in self._model_spec.model_id
-            or "gte-Qwen2" in self._model_spec.model_name
+            "gte" in self._model_spec.model_name.lower()
+            and "qwen2" in self._model_spec.model_name.lower()
         ):
             all_embeddings, all_token_nums = encode(
                 self._model,