查询向量与文档向量余弦相似度计算及排序代码修复问题

问题描述 #

我有一个queries字典，结构如下：

{'q1': ['similar',
  'law',
  'must',
  'obey',
  'construct',
  'aeroelast',
  'model',
  'heat',
  'high',
  'speed',
  'aircraft'],
 'q2': ['structur',
  'aeroelast',
  'problem',
  'associ',
  'flight',
  'high',
  'speed',
  'aircraft'],
 'q3': ['problem', 'heat', 'conduct', 'composit', 'slab', 'solv', 'far']
...
}

已通过代码将其转换为向量数组字典self.q_vecs，同时生成了文档向量self.doc_vecs，但编写的retrieve_n_rank_docs方法仅返回一个键值对，无法实现计算查询向量与文档向量的余弦相似度，按相似度降序排序文档ID的需求，期望输出结构：

{
    'q217': ['d983', 'd554', ..., 'd623'],
    'q99' : ['d716', 'd67', ..., 'd164'],
    ...
}

错误代码分析 #

现有retrieve_n_rank_docs方法存在以下问题：

• 每次循环查询时都重新初始化new_d = {}，导致之前的结果被覆盖
• 错误计算了查询向量之间的余弦相似度，而非需求的查询向量与文档向量的相似度
• 没有保存文档ID与向量的对应关系，无法将相似度结果映射到文档ID
• 方法没有返回最终生成的结果字典

修正后的完整代码 #

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import TruncatedSVD
from scipy.spatial import distance

class RetrievalSystem:
    def __init__(self, docs, num_concepts, min_df=1, alpha=1.0, beta=0.75, gamma=0.15):
        self.alpha, self.beta, self.gamma = alpha, beta, gamma
        
        # 保存文档ID的顺序，后续用于映射相似度结果
        self.doc_ids = list(docs.keys())
        
        # 创建文档-词项矩阵
        self.vec = TfidfVectorizer(tokenizer=str.split, min_df=min_df)
        doc_term_mat = self.vec.fit_transform([" ".join(docs[doc_id]) for doc_id in self.doc_ids])
        self.q_vecs = dict() # 存储查询向量
        
        # 用SVD降维
        self.svd = TruncatedSVD(n_components=num_concepts, random_state=42)
        self.doc_vecs = self.svd.fit_transform(doc_term_mat)

    def retrieve_n_rank_docs(self, queries, max_docs=-1):
        # 初始化结果字典
        ranked_docs = {}
        
        for query_id, query_terms in queries.items():
            # 转换查询为向量（如果未缓存）
            if query_id not in self.q_vecs:
                query_text = " ".join(query_terms)
                query_tfidf = self.vec.transform([query_text])
                query_vec = self.svd.transform(query_tfidf)
                self.q_vecs[query_id] = query_vec
            
            # 获取当前查询的向量
            q_vec = self.q_vecs[query_id].flatten()
            
            # 计算当前查询与所有文档的余弦相似度
            similarities = []
            for doc_vec in self.doc_vecs:
                sim = 1 - distance.cosine(q_vec, doc_vec)
                similarities.append(sim)
            
            # 将文档ID与相似度配对，按相似度降序排序
            doc_sim_pairs = list(zip(self.doc_ids, similarities))
            doc_sim_pairs.sort(key=lambda x: x[1], reverse=True)
            
            # 提取排序后的文档ID
            ranked_ids = [doc_id for doc_id, _ in doc_sim_pairs]
            
            # 限制返回的文档数量（如果max_docs有效）
            if max_docs > 0:
                ranked_ids = ranked_ids[:max_docs]
            
            # 存入结果字典
            ranked_docs[query_id] = ranked_ids
        
        # 返回最终结果
        return ranked_docs

关键修正点说明 #

1. 保存文档ID顺序：在__init__中添加self.doc_ids = list(docs.keys())，确保文档向量和文档ID一一对应
2. 结果字典初始化：将ranked_docs初始化在循环外部，避免每次循环覆盖之前的结果
3. 正确计算相似度：遍历所有文档向量，计算查询向量与每个文档向量的余弦相似度
4. 排序与结果提取：将文档ID和相似度配对后排序，提取排序后的文档ID列表
5. 返回结果：方法最后返回生成的结果字典，确保调用者能获取所有查询的排序结果

内容的提问来源于stack exchange，提问作者Hefe