Browse Source

Refa: Remove useless conver and fix a bug for DefaultRerank (#8887)

### What problem does this PR solve?

1. bug when re-try, we need to reset i.
2. remove useless convert

### Type of change

- [x] Refactoring
tags/v0.20.0
Stephen Hu 3 months ago
parent
commit
38b34116dd
No account linked to committer's email address
1 changed files with 5 additions and 3 deletions
  1. 5
    3
      rag/llm/rerank_model.py

+ 5
- 3
rag/llm/rerank_model.py View File

@@ -101,9 +101,10 @@ class DefaultRerank(Base):
old_dynamic_batch_size = self._dynamic_batch_size
if max_batch_size is not None:
self._dynamic_batch_size = max_batch_size
res = []
res = np.array([], dtype=float)
i = 0
while i < len(pairs):
cur_i = i
current_batch = self._dynamic_batch_size
max_retries = 5
retry_count = 0
@@ -111,7 +112,7 @@ class DefaultRerank(Base):
try:
# call subclass implemented batch processing calculation
batch_scores = self._compute_batch_scores(pairs[i : i + current_batch])
res.extend(batch_scores)
res = np.append(res, batch_scores)
i += current_batch
self._dynamic_batch_size = min(self._dynamic_batch_size * 2, 8)
break
@@ -119,6 +120,7 @@ class DefaultRerank(Base):
if "CUDA out of memory" in str(e) and current_batch > self._min_batch_size:
current_batch = max(current_batch // 2, self._min_batch_size)
self.torch_empty_cache()
i = cur_i # reset i to the start of the current batch
retry_count += 1
else:
raise
@@ -134,7 +136,7 @@ class DefaultRerank(Base):
scores = self._model.compute_score(batch_pairs)
else:
scores = self._model.compute_score(batch_pairs, max_length=max_length)
scores = sigmoid(np.array(scores)).tolist()
scores = sigmoid(np.array(scores))
if not isinstance(scores, Iterable):
scores = [scores]
return scores

Loading…
Cancel
Save