From cb6a4a9f6b253301729f63addf6a4e6da4f008d8 Mon Sep 17 00:00:00 2001
From: WNJXYK <wnjxyk@gmail.com>
Date: Sat, 27 May 2023 17:19:21 +0800
Subject: [PATCH] [FIX] Remove duplicate

---
 utils/knowledge.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/utils/knowledge.py b/utils/knowledge.py
index edb9734..f76f612 100644
--- a/utils/knowledge.py
+++ b/utils/knowledge.py
@@ -22,10 +22,14 @@ class Knowledge(object):
         embedding = self.knowledge.embedding_function(prompt)
         scores, indices = self.knowledge.index.search(np.array([embedding], dtype=np.float32), topk)
         docs = []
+        titles = set()
         for j, i in enumerate(indices[0]):
             if i == -1: continue
             if scores[0][j] > threshold: continue
-            docs.append(self.render_index(i, scores[0][j]))
+            item = self.render_index(i, scores[0][j])
+            if item["title"] in titles: continue
+            titles.add(item["title"])
+            docs.append(item)
         return docs
 
     def get_response(self, output: str) -> str:
@@ -37,4 +41,6 @@ class Knowledge(object):
         return res
     
 # knowledge = Knowledge()
-# print(knowledge.get_response(knowledge.query_prompt("酒后驾车")))
\ No newline at end of file
+# answer = knowledge.query_prompt("强奸男性犯法吗？")
+# print(answer)
+# print(knowledge.get_response(answer))
\ No newline at end of file