Browse Source
Fix token count
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
1 additions and
1 deletions
-
rerank.py
|
|
@ -67,7 +67,7 @@ class ReRank(NNOperator): |
|
|
texts[idx].append(text.strip()) |
|
|
texts[idx].append(text.strip()) |
|
|
|
|
|
|
|
|
tokenized = self.tokenizer(*texts, padding=True, truncation='longest_first', return_tensors="pt", max_length=self.max_length) |
|
|
tokenized = self.tokenizer(*texts, padding=True, truncation='longest_first', return_tensors="pt", max_length=self.max_length) |
|
|
token_count = torch.count_nonzero(tokenized['input_ids']) |
|
|
|
|
|
|
|
|
token_count = int(torch.count_nonzero(tokenized['input_ids'])) |
|
|
|
|
|
|
|
|
logits = self.model(**tokenized) |
|
|
logits = self.model(**tokenized) |
|
|
scores = self.post_proc(logits) |
|
|
scores = self.post_proc(logits) |
|
|
|
