From 3fd23dcc7da27fab78c11a5a17f41841e8a84ca6 Mon Sep 17 00:00:00 2001 From: zhangpu <1250681871@qq.com> Date: Mon, 13 Oct 2025 14:08:59 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E4=BA=86=E5=90=91=E9=87=8F?= =?UTF-8?q?=E5=8F=AF=E8=A7=86=E5=8C=96=E8=84=9A=E6=9C=AC=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- faiss_vector_db/visualize_embeddings.py | 214 ++++++++++++++++++++++++ 1 file changed, 214 insertions(+) create mode 100644 faiss_vector_db/visualize_embeddings.py diff --git a/faiss_vector_db/visualize_embeddings.py b/faiss_vector_db/visualize_embeddings.py new file mode 100644 index 0000000..9e58853 --- /dev/null +++ b/faiss_vector_db/visualize_embeddings.py @@ -0,0 +1,214 @@ +import argparse +import json +import os +import sys +from typing import Optional, Tuple + +import numpy as np +import matplotlib.pyplot as plt + +# 尝试可选依赖 +_umap_available = False +try: + import umap # type: ignore + _umap_available = True +except Exception: + _umap_available = False + +from sklearn.decomposition import PCA +from sklearn.manifold import TSNE +from sklearn.metrics import silhouette_score, pairwise_distances +from sklearn.preprocessing import StandardScaler + + +def load_embeddings_json(path: str) -> np.ndarray: + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + emb = np.asarray(data, dtype=np.float32) + if emb.ndim != 2: + raise ValueError(f"embeddings 不是二维数组: shape={emb.shape}") + return emb + + +def load_labels_json(path: str) -> np.ndarray: + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + labels = np.asarray(data, dtype=np.int64) + if labels.ndim != 1: + raise ValueError(f"labels 不是一维数组: shape={labels.shape}") + return labels + + +def subsample(X: np.ndarray, y: np.ndarray, max_points: Optional[int], seed: int = 42) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + n = X.shape[0] + idx = np.arange(n) + if max_points is not None and n > max_points: + rng = np.random.default_rng(seed) + sel = rng.choice(n, size=max_points, replace=False) + sel = np.sort(sel) + return X[sel], y[sel], sel + return X, y, idx + + +def diagnostics(X: np.ndarray, y: np.ndarray, reduced2d: Optional[np.ndarray] = None, method: str = "pca") -> str: + # 基本统计 + N, D = X.shape + means = X.mean(axis=0) + stds = X.std(axis=0) + mean_norm = float(np.linalg.norm(means)) + mean_std = float(stds.mean()) + zero_var_ratio = float((stds == 0).mean()) + + # 余弦相似度（抽样，避免 O(N^2)） + max_pairs = 20000 + rng = np.random.default_rng(0) + if N > 1: + sample_size = min(N, int(np.sqrt(max_pairs)) + 1) + idx = rng.choice(N, size=sample_size, replace=False) + Xs = X[idx] + # 归一化以计算余弦 + Xn = Xs / (np.linalg.norm(Xs, axis=1, keepdims=True) + 1e-12) + cos = Xn @ Xn.T + iu = np.triu_indices_from(cos, k=1) + cos_vals = cos[iu] + cos_mean = float(cos_vals.mean()) if cos_vals.size > 0 else float("nan") + cos_std = float(cos_vals.std()) if cos_vals.size > 0 else float("nan") + else: + cos_mean = float("nan") + cos_std = float("nan") + + # silhouette（类内/类间可分性），需要至少2类且每类>=2个点 + sil = float("nan") + unique, counts = np.unique(y, return_counts=True) + if unique.size >= 2 and (counts >= 2).all(): + # 为避免过大计算，对子集做 + subN = min(N, 5000) + if N > subN: + idx = rng.choice(N, size=subN, replace=False) + Xs, ys = X[idx], y[idx] + else: + Xs, ys = X, y + # 用欧氏距离 + try: + sil = float(silhouette_score(Xs, ys, metric="euclidean")) + except Exception: + sil = float("nan") + + # PCA前两维解释率（若提供reduced2d且method=pca，则更有意义） + pca_explained_2 = float("nan") + try: + pca = PCA(n_components=2, random_state=0) + pca.fit(X) + pca_explained_2 = float(pca.explained_variance_ratio_.sum()) + except Exception: + pass + + lines = [ + f"N={N}, D={D}, 类别数={unique.size}, 每类样本数范围=[{counts.min() if counts.size>0 else 0}, {counts.max() if counts.size>0 else 0}]", + f"均值向量范数={mean_norm:.6f}, 维度平均标准差={mean_std:.6f}, 零方差维度占比={zero_var_ratio*100:.2f}%", + f"余弦相似度: 均值={cos_mean:.6f}, 标准差={cos_std:.6f} (越接近1且方差小，越可能坍塌)", + f"Silhouette 分数={sil:.4f} (接近0或为负，类间不可分)", + f"PCA前两维累计解释率={pca_explained_2*100:.2f}%", + f"降维方法={method}", + ] + return "\n".join(lines) + + +def reduce_dim(X: np.ndarray, method: str, seed: int, tsne_perplexity: int, umap_n_neighbors: int, umap_min_dist: float) -> np.ndarray: + if method == "pca": + # 标准化后做PCA，通常更稳定 + Xs = StandardScaler(with_mean=True, with_std=True).fit_transform(X) + comp = PCA(n_components=2, random_state=seed) + Z = comp.fit_transform(Xs) + return Z + elif method == "tsne": + # t-SNE 对规模敏感，建议先PCA到50维再t-SNE + d = min(50, X.shape[1]) + Xp = PCA(n_components=d, random_state=seed).fit_transform(X) + ts = TSNE(n_components=2, perplexity=tsne_perplexity, learning_rate="auto", init="pca", random_state=seed, n_iter=1000) + Z = ts.fit_transform(Xp) + return Z + elif method == "umap": + if not _umap_available: + raise RuntimeError("未安装 umap-learn，请先: pip install umap-learn") + reducer = umap.UMAP(n_components=2, n_neighbors=umap_n_neighbors, min_dist=umap_min_dist, random_state=seed, metric="euclidean") + Z = reducer.fit_transform(X) + return Z + else: + raise ValueError(f"未知降维方法: {method}") + + +def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] = None) -> None: + plt.figure(figsize=(8, 7), dpi=120) + classes = np.unique(y) + # 构建调色板 + cmap = plt.get_cmap("tab20") + colors = [cmap(i % 20) for i in range(len(classes))] + + for i, c in enumerate(classes): + mask = (y == c) + plt.scatter(Z[mask, 0], Z[mask, 1], s=8, c=[colors[i]], label=str(c), alpha=0.7, edgecolors='none') + + plt.title(title) + plt.xlabel("Dim 1") + plt.ylabel("Dim 2") + if len(classes) <= 20: + plt.legend(markerscale=2, frameon=True, fontsize=8) + plt.tight_layout() + if out_path: + plt.savefig(out_path, bbox_inches="tight") + plt.show() + + +def main(): + parser = argparse.ArgumentParser(description="可视化高维 embedding 并进行坍塌诊断") + parser.add_argument("--embeddings", type=str, default=os.path.join( "faiss_index", "embeddings.json"), help="embeddings.json 路径") + parser.add_argument("--labels", type=str, default=os.path.join("faiss_index", "labels.json"), help="labels.json 路径") + parser.add_argument("--method", type=str, default="pca", choices=["pca", "tsne", "umap"], help="降维方法") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--max_points", type=int, default=None, help="抽样上限，避免t-SNE/UMAP过慢；None为全量") + parser.add_argument("--tsne_perplexity", type=int, default=30, help="t-SNE perplexity") + parser.add_argument("--umap_n_neighbors", type=int, default=15, help="UMAP n_neighbors") + parser.add_argument("--umap_min_dist", type=float, default=0.1, help="UMAP min_dist") + parser.add_argument("--no_center_scale", action="store_true", help="禁用标准化（默认PCA会标准化）") + args = parser.parse_args() + + if not os.path.exists(args.embeddings): + print(f"找不到 embeddings: {args.embeddings}", file=sys.stderr) + sys.exit(1) + if not os.path.exists(args.labels): + print(f"找不到 labels: {args.labels}", file=sys.stderr) + sys.exit(1) + + print("加载数据...") + X = load_embeddings_json(args.embeddings) + y = load_labels_json(args.labels) + + if X.shape[0] != y.shape[0]: + n = min(X.shape[0], y.shape[0]) + print(f"警告: embeddings与labels数量不一致, 截断到最小长度 {n}。", file=sys.stderr) + X = X[:n] + y = y[:n] + + # 抽样（t-SNE/UMAP 推荐抽样） + Xs, ys, sel_idx = subsample(X, y, args.max_points, seed=args.seed) + if Xs.shape[0] < X.shape[0]: + print(f"已抽样: {Xs.shape[0]}/{X.shape[0]}") + + # 降维 + print(f"降维方法: {args.method}") + Z = reduce_dim(Xs, args.method, args.seed, args.tsne_perplexity, args.umap_n_neighbors, args.umap_min_dist) + + # 诊断 + print("\n==== 诊断信息 ====") + print(diagnostics(Xs, ys, reduced2d=Z, method=args.method)) + + # 绘图 + out_png = os.path.join(os.path.dirname(args.embeddings), f"embedding_{args.method}_2d.png") + title = f"Embedding {args.method.upper()} 2D (N={Xs.shape[0]})" + plot_2d(Z, ys, title, out_path=out_png) + print(f"已保存图像: {out_png}") + + +if __name__ == "__main__": + main() \ No newline at end of file