Browse Source
Fix name bug
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
3 additions and
2 deletions
-
auto_transformers.py
|
|
|
@ -73,6 +73,7 @@ class AutoTransformers(NNOperator): |
|
|
|
else: |
|
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
self.model_name = self.map_model_names(model_name) |
|
|
|
print(1111, self.model_name) |
|
|
|
if tokenizer: |
|
|
|
self.tokenizer = tokenizer |
|
|
|
else: |
|
|
|
@ -252,8 +253,8 @@ class AutoTransformers(NNOperator): |
|
|
|
for line in data.split('\r\n'): |
|
|
|
line = line.replace(' ', '') |
|
|
|
if line.startswith('"name":'): |
|
|
|
name = line.split(':')[-1].replace('"', '').replace(',', '') |
|
|
|
default_sbert.append(name) |
|
|
|
n = line.split(':')[-1].replace('"', '').replace(',', '') |
|
|
|
default_sbert.append(n) |
|
|
|
if name in default_sbert: |
|
|
|
name = 'sentence-transformers/' + name |
|
|
|
return name |
|
|
|
|
