在 encode_texts ⽅法中,需要将输⼊⽂本列表转换为模型所需的 token 格式。请选择正确实现:
材料题 2:API ⽂档检索助⼿
API ⽂档检索助⼿是⼀类典型的检索增强系统。它先把 API ⽂档⽚段编码成向量并建⽴索引;当⽤户提出
问题时,再把问题编码成向量,与⽂档向量计算相似度,选出最相关的⽂档⽚段,⽤于后续回答⽣成。
补充材料:
1. tokenizer(texts, return_tensors="pt", padding=True, truncation=True) 会返回⼀个字 典,常⻅字段包括 input_ids 和 attention_mask。
2. 嵌⼊模型输出的 last_hidden_state 形状通常为 [batch_size, seq_len, hidden_dim]。
3. attention_mask 的形状通常为 [batch_size, seq_len],其中 1 表示真实 token,0 表示 padding token。
4. 对带 padding 的序列做平均池化时,应排除 padding 位置,否则短⽂本向量会被填充值稀释。
5. 若查询向量和⽂档向量都已做 L2 归⼀化,则⼆者点积等价于余弦相似度。
6. 如果查询向量形状为 [1, hidden_dim],⽂档矩阵形状为 [num_docs, hidden_dim],则⼀次性 计算所有⽂档相似度可使⽤矩阵乘法,结果形状为 [1, num_docs]。
请阅读代码,并根据描述完成空缺部分。
import torch
import torch.nn.functional as F
class DocumentRetriever:
def __init__(self, embedding_model, tokenizer, document_corpus):
self.embedding_model = embedding_model
self.tokenizer = tokenizer
self.document_corpus = document_corpus
self.doc_embeddings = self.index_documents(document_corpus)
def encode_texts(self, texts):
tokens = ____[1]____
tokens = {k: v.to(self.embedding_model.device) for k, v in tokens.items()}
with torch.no_grad():
model_output = self.embedding_model(**tokens)
hidden_states = model_output.last_hidden_state
mask = tokens["attention_mask"]
sentence_embeddings = ____[2]____
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
return sentence_embeddings
def index_documents(self, documents):
doc_vectors = self.encode_texts(documents)
return doc_vectors.cpu()
def retrieve_relevant_docs(self, query_embedding, top_k=3):
if query_embedding.dim() == 1:
query_embedding = query_embedding.unsqueeze(0)
query_normalized = F.normalize(query_embedding, p=2, dim=1).to(self.doc_embeddings.device)
similarity_scores = ____[3]____
similarity_scores = similarity_scores.squeeze(0)
top_scores, top_indices = torch.topk(
similarity_scores,
k=min(top_k, len(self.document_corpus)),
)
return [
{"document": self.document_corpus[idx], "relevance": float(top_scores[i])}
for i, idx in enumerate(top_indices)
]
self.tokenizer.encode(texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
self.embedding_model(texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
self.tokenizer.tokenize(texts, return_tensors="pt", padding=True, truncation=True)