Chapter 17
检索机器人
Notebooktransformers19 cells
检索机器人
Step1 读取faq数据
In [ ]python · cell 3
python
import pandas as pd
data = pd.read_csv("./law_faq.csv")
data.head()Step2 加载模型
In [ ]python · cell 5
python
from dual_model import DualModel
# 需要完成前置模型训练
dual_model = DualModel.from_pretrained("../12-sentence_similarity/dual_model/checkpoint-500/")
dual_model = dual_model.cuda()
dual_model.eval()
print("匹配模型加载成功！")In [ ]python · cell 6
python
from transformers import AutoTokenizer
tokenzier = AutoTokenizer.from_pretrained("hfl/chinese-macbert-base")
tokenzierStep3 将问题编码为向量
In [ ]python · cell 8
python
import torch
from tqdm import tqdm
questions = data["title"].to_list()
vectors = []
with torch.inference_mode():
for i in tqdm(range(0, len(questions), 32)):
batch_sens = questions[i: i + 32]
inputs = tokenzier(batch_sens, return_tensors="pt", padding=True, max_length=128, truncation=True)
inputs = {k: v.to(dual_model.device) for k, v in inputs.items()}
vector = dual_model.bert(**inputs)[1]
vectors.append(vector)
vectors = torch.concat(vectors, dim=0).cpu().numpy()
vectors.shapeStep4 创建索引
In [ ]python · cell 10
python
import faiss
index = faiss.IndexFlatIP(768)
faiss.normalize_L2(vectors)
index.add(vectors)
indexStep5 对问题进行向量编码
In [ ]python · cell 12
python
quesiton = "寻衅滋事"
with torch.inference_mode():
inputs = tokenzier(quesiton, return_tensors="pt", padding=True, max_length=128, truncation=True)
inputs = {k: v.to(dual_model.device) for k, v in inputs.items()}
vector = dual_model.bert(**inputs)[1]
q_vector = vector.cpu().numpy()
q_vector.shapeStep6 向量匹配(召回)
In [ ]python · cell 14
python
faiss.normalize_L2(q_vector)
scores, indexes = index.search(q_vector, 10)
topk_result = data.values[indexes[0].tolist()]
topk_result[:, 0]Step7 加载交互模型
In [ ]python · cell 16
python
from transformers import BertForSequenceClassification
# 需要完成前置模型训练
corss_model = BertForSequenceClassification.from_pretrained("../12-sentence_similarity/cross_model/checkpoint-500/")
corss_model = corss_model.cuda()
corss_model.eval()
print("模型加载成功！")Step8 最终预测(排序)
In [ ]python · cell 18
python
canidate = topk_result[:, 0].tolist()
ques = [quesiton] * len(canidate)
inputs = tokenzier(ques, canidate, return_tensors="pt", padding=True, max_length=128, truncation=True)
inputs = {k: v.to(corss_model.device) for k, v in inputs.items()}
with torch.inference_mode():
logits = corss_model(**inputs).logits.squeeze()
result = torch.argmax(logits, dim=-1)
resultIn [ ]python · cell 19
python
canidate_answer = topk_result[:, 1].tolist()
match_quesiton = canidate[result.item()]
final_answer = canidate_answer[result.item()]
match_quesiton, final_answer