pandas 的 'embedding' 列与文本 'embedding' 计算相似度

如何利用pandas或其它方式，算一个文本的embedding和dataframe中'embedding'字段的相似性，dataframe中还有'group_idx'字段，如何按相同的group_idx降序，取每个group_idx中前15个行数据输出

好问题，这个是向量检索 / 语义搜索里非常经典的一步 👍
我给你一个可直接用的 Pandas + NumPy 实战方案，再补充一点性能和工程上的小技巧。

---

一、问题拆解一下

你现在有：

• 一个查询文本 → 要算它的 embedding（记为 query_embedding）

• 一个 DataFrame，包含：

  • embedding：每行一个向量（list / np.array）
  • group_idx：分组字段

目标是：

1. 计算 query_embedding 与每一行 embedding 的相似度

2. 按 group_idx 分组

3. 在每个 group 内：

   • 按相似度 降序
   • 取 Top 15

4. 输出结果

---

二、推荐的相似度：Cosine Similarity

语义 embedding 基本都用 余弦相似度：

[
\text{cosine}(a, b) = \frac{a \cdot b}{|a| |b|}
]

---

三、基础实现（Pandas + NumPy）

1️⃣ 数据结构假设

import pandas as pd
import numpy as np

# df 示例
# df.columns = ['embedding', 'group_idx', ...]
# embedding 是 list 或 np.array，例如 [0.12, 0.33, ...]

---

2️⃣ 定义 cosine 相似度函数

def cosine_similarity(a, b):
    a = np.asarray(a)
    b = np.asarray(b)
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

---

3️⃣ 计算查询文本的 embedding

假设你已经算好了（比如用 OpenAI / sentence-transformers）：

query_embedding = np.array([...])  # shape = (dim,)

---

4️⃣ 计算每一行的相似度

df = df.copy()

df["similarity"] = df["embedding"].apply(
    lambda x: cosine_similarity(query_embedding, x)
)

---

5️⃣ 按 group_idx 分组，取每组 Top 15

这是关键一步 👇

result = (
    df
    .sort_values("similarity", ascending=False)
    .groupby("group_idx", group_keys=False)
    .head(15)
)

✅ 得到的 result 就是：

• 每个 group_idx
• 相似度降序
• 前 15 条

---

四、性能更好的向量化版本（强烈推荐）

如果数据量大（比如几万 / 几十万行），不要用 apply。

1️⃣ 先把 embedding 堆成矩阵

emb_matrix = np.vstack(df["embedding"].values)  # shape: (N, dim)

---

2️⃣ 一次性算 cosine similarity

# 归一化
emb_norm = emb_matrix / np.linalg.norm(emb_matrix, axis=1, keepdims=True)
query_norm = query_embedding / np.linalg.norm(query_embedding)

# 向量点积
similarities = emb_norm @ query_norm

---

3️⃣ 写回 DataFrame

df["similarity"] = similarities

---

4️⃣ 分组 Top 15（同上）

result = (
    df
    .sort_values("similarity", ascending=False)
    .groupby("group_idx", group_keys=False)
    .head(15)
)

---

五、如果你用的是 sentence-transformers / FAISS

✔ sentence-transformers 自带 cosine

from sentence_transformers.util import cos_sim

scores = cos_sim(query_embedding, emb_matrix)[0].cpu().numpy()
df["similarity"] = scores

---

✔ 如果数据再大一点（百万级）

建议直接：

• 按 group_idx 分多个 FAISS index

• 或者：

  • 先粗召回 Top K
  • 再在 group 内精排

我也可以帮你设计这个结构。

---

六、常见坑提醒（很重要）

1. embedding 一定要是同一模型算的

2. cosine 前记得归一化

3. groupby().head() 前必须先 sort_values

4. embedding 存数据库时建议：

   • list → numpy array
   • 或者直接 float32 matrix

---

七、如果你愿意，我可以继续帮你：

• 🔹 用 FAISS 按 group_idx 做分组检索
• 🔹 把这套逻辑封装成一个函数 / 类
• 🔹 针对你的 embedding 模型（OpenAI / bge / e5）给最优写法
• 🔹 直接帮你优化到生产级速度

你可以直接贴一段你的 df.head() 或 embedding 维度，我就能更精准地改。