From 77363a087562fa2ff27109b60930a419ef15640f Mon Sep 17 00:00:00 2001
From: KevinHuSh <kevinhu.sh@gmail.com>
Date: Thu, 30 May 2024 12:55:17 +0800
Subject: [PATCH] fix bge rerank normalize issue (#988)

### What problem does this PR solve?


### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 rag/llm/rerank_model.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py
index c1c12f113..5039e11a0 100644
--- a/rag/llm/rerank_model.py
+++ b/rag/llm/rerank_model.py
@@ -24,6 +24,8 @@ import numpy as np
 from api.utils.file_utils import get_home_cache_dir
 from rag.utils import num_tokens_from_string, truncate
 
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
 
 class Base(ABC):
     def __init__(self, key, model_name):
@@ -69,6 +71,7 @@ class DefaultRerank(Base):
         res = []
         for i in range(0, len(pairs), batch_size):
             scores = self._model.compute_score(pairs[i:i + batch_size], max_length=2048)
+            scores = sigmoid(np.array(scores))
             res.extend(scores)
         return np.array(res), token_count