fixed bug no set max_length and special token

LazyAGI · Sep 11, 2024 · b1d1bb0 · b1d1bb0
1 parent 30bd933
commit b1d1bb0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lazyllm/components/embedding/embed.py b/lazyllm/components/embedding/embed.py
@@ -26,7 +26,8 @@ def load_embed(self):
 
     def __call__(self, string):
         lazyllm.call_once(self.init_flag, self.load_embed)
-        encoded_input = self.tokenizer(string, padding=True, truncation=True, return_tensors='pt').to(self.device)
+        encoded_input = self.tokenizer(string, padding=True, truncation=True, return_tensors='pt',
+                                       max_length=512, add_special_tokens=True).to(self.device)
         with torch.no_grad():
             model_output = self.embed(**encoded_input)
             sentence_embeddings = model_output[0][:, 0]