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