Browse Source
Output token usage
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
2 additions and
1 deletions
-
rerank.py
|
|
|
@ -67,6 +67,7 @@ class ReRank(NNOperator): |
|
|
|
texts[idx].append(text.strip()) |
|
|
|
|
|
|
|
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']) |
|
|
|
|
|
|
|
logits = self.model(**tokenized) |
|
|
|
scores = self.post_proc(logits) |
|
|
|
@ -78,7 +79,7 @@ class ReRank(NNOperator): |
|
|
|
else: |
|
|
|
re_docs = [docs[i] for i in re_ids if scores[i] >= self._threshold] |
|
|
|
re_scores = [scores[i] for i in re_ids if scores[i] >= self._threshold] |
|
|
|
return re_docs, re_scores |
|
|
|
return re_docs, re_scores, token_count |
|
|
|
|
|
|
|
|
|
|
|
def post_proc(self, logits): |
|
|
|
|
