Browse Source
Fix token
Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
main
1 changed files with
12 additions and
15 deletions
-
s_bert.py
|
|
|
@ -20,6 +20,7 @@ from functools import partial |
|
|
|
|
|
|
|
import torch |
|
|
|
from sentence_transformers import SentenceTransformer |
|
|
|
from transformers import AutoTokenizer, AutoConfig |
|
|
|
|
|
|
|
from towhee.operator import NNOperator |
|
|
|
try: |
|
|
|
@ -105,10 +106,6 @@ class STransformers(NNOperator): |
|
|
|
return ['onnx'] |
|
|
|
|
|
|
|
def get_tokenizer(self): |
|
|
|
if hasattr(self._model, "tokenize"): |
|
|
|
return self._model.tokenize |
|
|
|
else: |
|
|
|
from transformers import AutoTokenizer, AutoConfig |
|
|
|
try: |
|
|
|
tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/' + self.model_name) |
|
|
|
conf = AutoConfig.from_pretrained('sentence-transformers/' + self.model_name) |
|
|
|
|
