From 4b3eeaa6efc479030bcf655fb633f7fedd3dd75a Mon Sep 17 00:00:00 2001
From: shizzgar <39241841+shizzgar@users.noreply.github.com>
Date: Mon, 18 Nov 2024 07:05:52 +0300
Subject: [PATCH] Added LocalAI support for rerank models (#3446)

### What problem does this PR solve?

Hi there!
LocalAI added support of rerank models
https://localai.io/features/reranker/

I've implemented LocalAIRerank class (typically copied it from
OpenAI_APIRerank class).
Also, LocalAI model response with 500 error code if len of "documents"
is less than 2 in similarity check.
So I've added the second "document" on RERANK model connection check in
`api/apps/llm_app.py`.

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
---
 api/apps/llm_app.py     |  2 +-
 rag/llm/__init__.py     |  1 +
 rag/llm/rerank_model.py | 39 +++++++++++++++++++++++++++++++++++++--
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/api/apps/llm_app.py b/api/apps/llm_app.py
index 9b5901765..75ad45a3c 100644
--- a/api/apps/llm_app.py
+++ b/api/apps/llm_app.py
@@ -238,7 +238,7 @@ def add_llm():
             base_url=llm["api_base"]
         )
         try:
-            arr, tc = mdl.similarity("Hello~ Ragflower!", ["Hi, there!"])
+            arr, tc = mdl.similarity("Hello~ Ragflower!", ["Hi, there!", "Ohh, my friend!"])
             if len(arr) == 0 or tc == 0:
                 raise Exception("Not known.")
         except Exception as e:
diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py
index 1a28a70e1..6d26a1fe6 100644
--- a/rag/llm/__init__.py
+++ b/rag/llm/__init__.py
@@ -110,6 +110,7 @@ ChatModel = {
 }
 
 RerankModel = {
+    "LocalAI":LocalAIRerank,
     "BAAI": DefaultRerank,
     "Jina": JinaRerank,
     "Youdao": YoudaoRerank,
diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py
index cac7b342d..58ea0decb 100644
--- a/rag/llm/rerank_model.py
+++ b/rag/llm/rerank_model.py
@@ -185,11 +185,46 @@ class XInferenceRerank(Base):
 
 class LocalAIRerank(Base):
     def __init__(self, key, model_name, base_url):
-        pass
+        if base_url.find("/rerank") == -1:
+            self.base_url = urljoin(base_url, "/rerank")
+        else:
+            self.base_url = base_url
+        self.headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {key}"
+        }
+        self.model_name = model_name.replace("___LocalAI","")
 
     def similarity(self, query: str, texts: list):
-        raise NotImplementedError("The LocalAIRerank has not been implement")
+        # noway to config Ragflow , use fix setting
+        texts = [truncate(t, 500) for t in texts]
+        data = {
+            "model": self.model_name,
+            "query": query,
+            "documents": texts,
+            "top_n": len(texts),
+        }
+        token_count = 0
+        for t in texts:
+            token_count += num_tokens_from_string(t)
+        res = requests.post(self.base_url, headers=self.headers, json=data).json()
+        rank = np.zeros(len(texts), dtype=float)
+        if 'results' not in res:
+            raise ValueError("response not contains results\n" + str(res))
+        for d in res["results"]:
+            rank[d["index"]] = d["relevance_score"]
 
+        # Normalize the rank values to the range 0 to 1
+        min_rank = np.min(rank)
+        max_rank = np.max(rank)
+
+        # Avoid division by zero if all ranks are identical
+        if max_rank - min_rank != 0:
+            rank = (rank - min_rank) / (max_rank - min_rank)
+        else:
+            rank = np.zeros_like(rank)
+
+        return rank, token_count
 
 class NvidiaRerank(Base):
     def __init__(