Upload ruler_tsne.py with huggingface_hub

ruler_tsne.py

@@ -0,0 +1,215 @@
+#!/usr/bin/env python3
+"""把输入文本 + ruler 200 条样本一起送进 Qwen3-Embedding-8B，做 t-SNE 可视化。
+
+用法：
+  # 输入文本直接给字符串
+  python ruler_tsne.py --input-text "user_0: hi\nuser_1: hello"
+
+  # 输入文本从文件读
+  python ruler_tsne.py --input-text /path/to/conv.txt
+
+  # 自定义路径 / 输出 / 批大小 / 最大长度
+  python ruler_tsne.py \
+    --input-text /path/to/conv.txt \
+    --ruler /mnt/.../ruler_items.json \
+    --model /mnt/.../Qwen3-Embedding-8B \
+    --output ruler_tsne.png \
+    --max-length 4096 \
+    --batch-size 4 \
+    --perplexity 20
+
+输出：
+  - ruler_tsne.png  -- 二维散点图，200 个 ruler 点按 score 着色（红=严重→绿=不严重），
+    每个点上标 rank 编号；输入文本用红色五角星标 INPUT。
+  - 控制台同时打印 top-5 最相似的 ruler items（按 cosine 相似度）。
+"""
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from sklearn.manifold import TSNE
+
+
+DEFAULT_MODEL = "/mnt/bn/tns-algo-ue-my/biaowu/WorkSpace/Models/Qwen3-Embedding-8B"
+DEFAULT_RULER = "/mnt/bn/tns-algo-ue-my/biaowu/aipf_dm_metric/ranking_moderation/data/dm/youth_sexual_and_physical_abuse_aigt_v009/ranking_bucket/ruler_items.json"
+
+
+def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    sequence_lengths = attention_mask.sum(dim=1) - 1
+    bsz = last_hidden_states.shape[0]
+    return last_hidden_states[torch.arange(bsz, device=last_hidden_states.device), sequence_lengths]
+
+
+@torch.no_grad()
+def encode(texts, tokenizer, model, max_length=4096, batch_size=4):
+    embs = []
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i:i + batch_size]
+        d = tokenizer(batch, padding=True, truncation=True,
+                      max_length=max_length, return_tensors="pt").to(model.device)
+        out = model(**d)
+        e = last_token_pool(out.last_hidden_state, d["attention_mask"])
+        e = F.normalize(e, p=2, dim=1)
+        embs.append(e.cpu().float())
+        del out, d, e
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print(f"  encoded {min(i + batch_size, len(texts))}/{len(texts)}", flush=True)
+    return torch.cat(embs, dim=0).numpy()
+
+
+def load_ruler_items(path: str):
+    with open(path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+    if isinstance(data, list):
+        items = data
+    else:
+        for k in ("items", "ruler_items", "data"):
+            if k in data and isinstance(data[k], list):
+                items = data[k]
+                break
+        else:
+            raise ValueError("unexpected ruler json structure")
+    out = []
+    for it in items:
+        inner = it.get("item", {}) if isinstance(it.get("item"), dict) else {}
+        conv = inner.get("conv_text") or it.get("conv_text") or ""
+        out.append({
+            "rank": it.get("rank"),
+            "score": it.get("score"),
+            "item_id": it.get("item_id"),
+            "text": conv,
+        })
+    return out
+
+
+def resolve_input(arg: str) -> str:
+    if arg == "-":
+        return sys.stdin.read().strip()
+    p = Path(arg)
+    if p.exists() and p.is_file():
+        return p.read_text(encoding="utf-8").strip()
+    return arg
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--input-text", required=True,
+                   help="原始文本字符串、文件路径，或 '-' 表示从 stdin 读")
+    p.add_argument("--ruler", default=DEFAULT_RULER)
+    p.add_argument("--model", default=DEFAULT_MODEL)
+    p.add_argument("--output", default="ruler_tsne.png")
+    p.add_argument("--max-length", type=int, default=4096)
+    p.add_argument("--batch-size", type=int, default=4)
+    p.add_argument("--perplexity", type=float, default=20.0)
+    p.add_argument("--label-fontsize", type=float, default=5,
+                   help="rank 编号的字号，太挤就调小")
+    p.add_argument("--cpu", action="store_true", help="强制走 CPU（不推荐，巨慢）")
+    p.add_argument("--no-flash-attn", action="store_true",
+                   help="不用 flash-attn-2（环境没装就加这个）")
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # ---- 1) 加载尺子 ----
+    print(f"[1/4] 读 ruler: {args.ruler}")
+    items = load_ruler_items(args.ruler)
+    print(f"      -> {len(items)} ruler items")
+
+    input_text = resolve_input(args.input_text)
+    print(f"      input text length: {len(input_text)} chars")
+    texts = [input_text] + [it["text"] for it in items]
+
+    # ---- 2) 加载模型 ----
+    print(f"[2/4] 加载模型: {args.model}")
+    device = "cpu" if args.cpu else ("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"      device: {device}")
+    model_kwargs = {}
+    if device == "cuda":
+        model_kwargs["torch_dtype"] = torch.float16
+        if not args.no_flash_attn:
+            try:
+                model_kwargs["attn_implementation"] = "flash_attention_2"
+            except Exception:
+                pass
+    tokenizer = AutoTokenizer.from_pretrained(args.model, padding_side="left")
+    model = AutoModel.from_pretrained(args.model, **model_kwargs).to(device).eval()
+
+    # ---- 3) 编码 ----
+    print(f"[3/4] 编码 {len(texts)} 条（batch_size={args.batch_size}, max_length={args.max_length}）")
+    embeddings = encode(texts, tokenizer, model,
+                        max_length=args.max_length, batch_size=args.batch_size)
+    print(f"      embeddings: {embeddings.shape}")
+
+    # ---- 4) t-SNE + 画图 ----
+    print(f"[4/4] t-SNE (perplexity={args.perplexity}) + 画图")
+    tsne = TSNE(n_components=2, perplexity=args.perplexity,
+                init="pca", random_state=42, metric="cosine")
+    xy = tsne.fit_transform(embeddings)
+
+    input_xy = xy[0]
+    ruler_xy = xy[1:]
+    ranks = np.array([it["rank"] for it in items])
+    scores = np.array([it["score"] for it in items], dtype=float)
+
+    fig, ax = plt.subplots(figsize=(14, 12), dpi=130)
+    sc = ax.scatter(
+        ruler_xy[:, 0], ruler_xy[:, 1],
+        c=scores, cmap="RdYlGn_r",
+        s=45, alpha=0.85,
+        edgecolor="black", linewidth=0.3,
+    )
+    cbar = plt.colorbar(sc, ax=ax, shrink=0.8)
+    cbar.set_label("ruler score (high = more severe)")
+
+    # 标 rank 编号
+    for (x, y), r in zip(ruler_xy, ranks):
+        ax.annotate(str(r), (x, y),
+                    fontsize=args.label_fontsize,
+                    ha="center", va="center",
+                    alpha=0.85)
+
+    # 输入点
+    ax.scatter([input_xy[0]], [input_xy[1]],
+               marker="*", s=750, c="red",
+               edgecolor="black", linewidth=1.5,
+               zorder=10, label="INPUT")
+    ax.annotate("INPUT", input_xy,
+                fontsize=12, fontweight="bold", color="red",
+                xytext=(10, 10), textcoords="offset points")
+
+    ax.set_title("t-SNE: input + 200 ruler items (Qwen3-Embedding-8B)")
+    ax.set_xlabel("t-SNE 1")
+    ax.set_ylabel("t-SNE 2")
+    ax.legend(loc="best")
+    plt.tight_layout()
+    plt.savefig(args.output, dpi=130, bbox_inches="tight")
+    print(f"      saved: {args.output}")
+
+    # ---- 5) Top-5 最相似 ruler items ----
+    sims = embeddings[1:] @ embeddings[0]      # cosine since L2 normalized
+    top5 = np.argsort(-sims)[:5]
+    print("\nTop-5 nearest ruler items by cosine similarity:")
+    for idx in top5:
+        it = items[idx]
+        print(f"  rank={it['rank']:>3}  score={it['score']:.2f}  "
+              f"sim={sims[idx]:.4f}  id={it['item_id']}")
+
+
+if __name__ == "__main__":
+    main()