logo
Browse Source

Fix inputs

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 2 years ago
parent
commit
92d43c7aa0
  1. 2
      codebert.py

2
codebert.py

@ -77,7 +77,7 @@ class CodeBert(NNOperator):
tokens = [self.tokenizer.cls_token, '<encoder-only>', self.tokenizer.sep_token] + tokens + \ tokens = [self.tokenizer.cls_token, '<encoder-only>', self.tokenizer.sep_token] + tokens + \
[self.tokenizer.sep_token] [self.tokenizer.sep_token]
tokens_ids = self.tokenizer.convert_tokens_to_ids(tokens) tokens_ids = self.tokenizer.convert_tokens_to_ids(tokens)
inputs = torch.tensor(tokens_ids).to(self.device)
inputs = torch.tensor(tokens_ids).unsqueeze(0).to(self.device)
except Exception as e: except Exception as e:
log.error(f'Invalid input for the tokenizer: {self.model_name}') log.error(f'Invalid input for the tokenizer: {self.model_name}')
raise e raise e

Loading…
Cancel
Save