[FIX] Remove duplicate

This commit is contained in:
WNJXYK
2023-05-27 17:19:21 +08:00
parent dc849bd282
commit cb6a4a9f6b

View File

@ -22,10 +22,14 @@ class Knowledge(object):
embedding = self.knowledge.embedding_function(prompt) embedding = self.knowledge.embedding_function(prompt)
scores, indices = self.knowledge.index.search(np.array([embedding], dtype=np.float32), topk) scores, indices = self.knowledge.index.search(np.array([embedding], dtype=np.float32), topk)
docs = [] docs = []
titles = set()
for j, i in enumerate(indices[0]): for j, i in enumerate(indices[0]):
if i == -1: continue if i == -1: continue
if scores[0][j] > threshold: 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 return docs
def get_response(self, output: str) -> str: def get_response(self, output: str) -> str:
@ -37,4 +41,6 @@ class Knowledge(object):
return res return res
# knowledge = Knowledge() # knowledge = Knowledge()
# print(knowledge.get_response(knowledge.query_prompt("酒后驾车"))) # answer = knowledge.query_prompt("强奸男性犯法吗?")
# print(answer)
# print(knowledge.get_response(answer))