Truncate text for zhipu embedding.

infiniflow · Jan 15, 2025 · 8426b98 · 8426b98
1 parent 7944aac
commit 8426b98
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py
@@ -217,6 +217,14 @@ def __init__(self, key, model_name="embedding-2", **kwargs):
     def encode(self, texts: list):
         arr = []
         tks_num = 0
+        MAX_LEN = -1
+        if self.model_name.lower() == "embedding-2":
+            MAX_LEN = 512
+        if self.model_name.lower() == "embedding-3":
+            MAX_LEN = 3072
+        if MAX_LEN > 0:
+            texts = [truncate(t, MAX_LEN) for t in texts]
+
         for txt in texts:
             res = self.client.embeddings.create(input=txt,
                                                 model=self.model_name)