68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
|
import os
|
|||
|
|
import json
|
|||
|
|
import pickle
|
|||
|
|
import numpy as np
|
|||
|
|
from typing import Tuple
|
|||
|
|
from sentence_transformers import SentenceTransformer
|
|||
|
|
|
|||
|
|
from config.config import knowledge_json_path, knowledge_pkl_path
|
|||
|
|
from util.encode import load_embedding, encode_qa
|
|||
|
|
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
读取知识库
|
|||
|
|
"""
|
|||
|
|
def load_knowledge() -> Tuple[list, list]:
|
|||
|
|
# 如果 pkl 不存在，则先编码存储
|
|||
|
|
if not os.path.exists(knowledge_pkl_path):
|
|||
|
|
encode_qa(knowledge_json_path, knowledge_pkl_path)
|
|||
|
|
|
|||
|
|
# 加载 json 和 pkl
|
|||
|
|
with open(knowledge_json_path, 'r', encoding='utf-8') as f1, open(knowledge_pkl_path, 'rb') as f2:
|
|||
|
|
knowledge = json.load(f1)
|
|||
|
|
encoded_knowledge = pickle.load(f2)
|
|||
|
|
return knowledge, encoded_knowledge
|
|||
|
|
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
召回 top_k 个相关的文本段
|
|||
|
|
"""
|
|||
|
|
def find_top_k(
|
|||
|
|
emb: SentenceTransformer,
|
|||
|
|
query: str,
|
|||
|
|
knowledge: list,
|
|||
|
|
encoded_knowledge: list,
|
|||
|
|
k=3
|
|||
|
|
) -> list[str]:
|
|||
|
|
# 编码 query
|
|||
|
|
query_embedding = emb.encode(query)
|
|||
|
|
|
|||
|
|
# 查找 top_k
|
|||
|
|
scores = query_embedding @ encoded_knowledge.T
|
|||
|
|
# 使用 argpartition 找出每行第 k 个大的值的索引，第 k 个位置左侧都是比它大的值，右侧都是比它小的值
|
|||
|
|
top_k_indices = np.argpartition(scores, -k)[-k:]
|
|||
|
|
# 由于 argpartition 不保证顺序，我们需要对提取出的 k 个索引进行排序
|
|||
|
|
top_k_values_sorted_indices = np.argsort(scores[top_k_indices])[::-1]
|
|||
|
|
top_k_indices = top_k_indices[top_k_values_sorted_indices]
|
|||
|
|
|
|||
|
|
# 返回
|
|||
|
|
contents = [knowledge[index] for index in top_k_indices]
|
|||
|
|
return contents
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
emb = load_embedding()
|
|||
|
|
knowledge, encoded_knowledge = load_knowledge()
|
|||
|
|
query = "认知心理学研究哪些心理活动？"
|
|||
|
|
contents = find_top_k(emb, query, knowledge, encoded_knowledge, 2)
|
|||
|
|
print('召回的 top-k 条相关内容如下：')
|
|||
|
|
print(json.dumps(contents, ensure_ascii=False, indent=2))
|
|||
|
|
# 这里我没实现 LLM 部分，如果有 LLM
|
|||
|
|
## 1. 读取 LLM
|
|||
|
|
## 2. 将 contents 拼接为 prompt，传给 LLM，作为 {已知内容}
|
|||
|
|
## 3. 要求 LLM 根据已知内容回复
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|