如何将余弦相似度结果转换为指定结构的Pandas DataFrame

你可以通过先收集所有结果行数据，再一次性生成DataFrame的方式实现需求，修改后的代码如下：

import pandas as pd
# 假设model、util、queries、corpus、corpus_embeddings已提前定义

# 初始化空列表存储结果数据
results = []
top_k = min(5, len(corpus))

for query in queries:
    query_embedding = model.encode(query, convert_to_tensor=True)
    cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]
    top_results = torch.topk(cos_scores, k=top_k)

    # 遍历top结果，将数据存入列表
    for score, idx in zip(top_results[0], top_results[1]):
        # 每个元素是(查询语句, 相似语料句子, 相似度分数)的元组
        results.append((query, corpus[idx], score.item()))

# 转换为指定结构的DataFrame
df1 = pd.DataFrame(results, columns=['Query', 'Corpus', 'Similarity Score'])
# 如果不需要分数列，可改为：
# df1 = pd.DataFrame(results, columns=['Query', 'Corpus'])

print(df1)

关键说明： #

• 用空列表results逐行存储每个查询对应的所有top5结果，避免在循环中频繁修改DataFrame（这种方式效率更高）
• score.item()用于将PyTorch张量转换为Python原生数值，方便存入DataFrame
• 最后通过pd.DataFrame直接将列表转换为DataFrame，指定列名即可匹配你需要的结构

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