需要对序列隐藏状态做 masked mean pooling，得到句⼦级向量。请选择正确实现：

材料题 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)

       ]