valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit

Browse Source 
Signed-off-by: Ilan Bigio <ilan@openai.com>
This commit is contained in:
Ilan Bigio
2025-04-16 12:56:08 -04:00
commit 59a180ddec
163 changed files with 30587 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

547
codex-cli/examples/prompt-analyzer/template/cluster_prompts.py Normal file
View File

@@ -0,0 +1,547 @@
#!/usr/bin/env python3
"""End‑to‑end pipeline for analysing a collection of text prompts.
The script performs the following steps:
1. Read a CSV file that must contain a column named ``prompt``. If an
``act`` column is present it is used purely for reporting purposes.
2. Create embeddings via the OpenAI API (``text-embedding-3-small`` by
default). The user can optionally provide a JSON cache path so the
expensive embedding step is only executed for new / unseen texts.
3. Cluster the resulting vectors either with K‑Means (automatically picking
*k* through the silhouette score) or with DBSCAN. Outliers are flagged
as cluster ``-1`` when DBSCAN is selected.
4. Ask a Chat Completion model (``gpt-4o-mini`` by default) to come up with a
short name and description for every cluster.
5. Write a human‑readable Markdown report (default: ``analysis.md``).
6. Generate a couple of diagnostic plots (cluster sizes and a t‑SNE scatter
plot) and store them in ``plots/``.
The script is intentionally opinionated yet configurable via a handful of CLI
options – run ``python cluster_prompts.py --help`` for details.
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from typing import Any, Sequence
import numpy as np
import pandas as pd
# External, heavy‑weight libraries are imported lazily so that users running the
# ``--help`` command do not pay the startup cost.
def parse_cli() -> argparse.Namespace: # noqa: D401
"""Parse command‑line arguments."""
parser = argparse.ArgumentParser(
prog="cluster_prompts.py",
description="Embed, cluster and analyse text prompts via the OpenAI API.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("--csv", type=Path, default=Path("prompts.csv"), help="Input CSV file.")
parser.add_argument(
"--cache",
type=Path,
default=None,
help="Optional JSON cache for embeddings (will be created if it does not exist).",
)
parser.add_argument(
"--embedding-model",
default="text-embedding-3-small",
help="OpenAI embedding model to use.",
)
parser.add_argument(
"--chat-model",
default="gpt-4o-mini",
help="OpenAI chat model for cluster descriptions.",
)
# Clustering parameters
parser.add_argument(
"--cluster-method",
choices=["kmeans", "dbscan"],
default="kmeans",
help="Clustering algorithm to use.",
)
parser.add_argument(
"--k-max",
type=int,
default=10,
help="Upper bound for k when the kmeans method is selected.",
)
parser.add_argument(
"--dbscan-min-samples",
type=int,
default=3,
help="min_samples parameter for DBSCAN (only relevant when dbscan is selected).",
)
# Output paths
parser.add_argument(
"--output-md", type=Path, default=Path("analysis.md"), help="Markdown report path."
)
parser.add_argument(
"--plots-dir", type=Path, default=Path("plots"), help="Directory that will hold PNG plots."
)
return parser.parse_args()
# ---------------------------------------------------------------------------
# Embedding helpers
# ---------------------------------------------------------------------------
def _lazy_import_openai(): # noqa: D401
"""Import *openai* only when needed to keep startup lightweight."""
try:
import openai # type: ignore
return openai
except ImportError as exc: # pragma: no cover – we do not test missing deps.
raise SystemExit(
"The 'openai' package is required but not installed.\n"
"Run 'pip install openai' and try again."
) from exc
def embed_texts(texts: Sequence[str], model: str, batch_size: int = 100) -> list[list[float]]:
"""Embed *texts* with OpenAI and return a list of vectors.
Uses batching for efficiency but remains on the safe side regarding current
OpenAI rate limits (can be adjusted by changing *batch_size*).
"""
openai = _lazy_import_openai()
client = openai.OpenAI()
embeddings: list[list[float]] = []
for batch_start in range(0, len(texts), batch_size):
batch = texts[batch_start : batch_start + batch_size]
response = client.embeddings.create(input=batch, model=model)
# The API returns the vectors in the same order as the input list.
embeddings.extend(data.embedding for data in response.data)
return embeddings
def load_or_create_embeddings(
prompts: pd.Series, *, cache_path: Path | None, model: str
) -> pd.DataFrame:
"""Return a *DataFrame* with one row per prompt and the embedding columns.
* If *cache_path* is provided and exists, known embeddings are loaded from
the JSON cache so they don't have to be re‑generated.
* Missing embeddings are requested from the OpenAI API and subsequently
appended to the cache.
* The returned DataFrame has the same index as *prompts*.
"""
cache: dict[str, list[float]] = {}
if cache_path and cache_path.exists():
try:
cache = json.loads(cache_path.read_text())
except json.JSONDecodeError: # pragma: no cover – unlikely.
print("⚠️ Cache file exists but is not valid JSON – ignoring.", file=sys.stderr)
missing_mask = ~prompts.isin(cache)
if missing_mask.any():
texts_to_embed = prompts[missing_mask].tolist()
print(f"Embedding {len(texts_to_embed)} new prompt(s)…", flush=True)
new_embeddings = embed_texts(texts_to_embed, model=model)
# Update cache (regardless of whether we persist it to disk later on).
cache.update(dict(zip(texts_to_embed, new_embeddings)))
if cache_path:
cache_path.parent.mkdir(parents=True, exist_ok=True)
cache_path.write_text(json.dumps(cache))
# Build a consistent embeddings matrix
vectors = prompts.map(cache.__getitem__).tolist() # type: ignore[arg-type]
mat = np.array(vectors, dtype=np.float32)
return pd.DataFrame(mat, index=prompts.index)
# ---------------------------------------------------------------------------
# Clustering helpers
# ---------------------------------------------------------------------------
def _lazy_import_sklearn_cluster():
"""Lazy import helper for scikit‑learn *cluster* sub‑module."""
# Importing scikit‑learn is slow; defer until needed.
from sklearn.cluster import DBSCAN, KMeans # type: ignore
from sklearn.metrics import silhouette_score # type: ignore
from sklearn.preprocessing import StandardScaler # type: ignore
return KMeans, DBSCAN, silhouette_score, StandardScaler
def cluster_kmeans(matrix: np.ndarray, k_max: int) -> np.ndarray:
"""Auto‑select *k* (in ``[2, k_max]``) via Silhouette score and cluster."""
KMeans, _, silhouette_score, _ = _lazy_import_sklearn_cluster()
best_k = None
best_score = -1.0
best_labels: np.ndarray | None = None
for k in range(2, k_max + 1):
model = KMeans(n_clusters=k, random_state=42, n_init="auto")
labels = model.fit_predict(matrix)
try:
score = silhouette_score(matrix, labels)
except ValueError:
# Occurs when a cluster ended up with 1 sample – skip.
continue
if score > best_score:
best_k = k
best_score = score
best_labels = labels
if best_labels is None: # pragma: no cover – highly unlikely.
raise RuntimeError("Unable to find a suitable number of clusters.")
print(f"K‑Means selected k={best_k} (silhouette={best_score:.3f}).", flush=True)
return best_labels
def cluster_dbscan(matrix: np.ndarray, min_samples: int) -> np.ndarray:
"""Cluster with DBSCAN; *eps* is estimated via the k‑distance method."""
_, DBSCAN, _, StandardScaler = _lazy_import_sklearn_cluster()
# Scale features – DBSCAN is sensitive to feature scale.
scaler = StandardScaler()
matrix_scaled = scaler.fit_transform(matrix)
# Heuristic: use the median of the distances to the ``min_samples``‑th
# nearest neighbour as eps. This is a commonly used rule of thumb.
from sklearn.neighbors import NearestNeighbors # type: ignore # lazy import
neigh = NearestNeighbors(n_neighbors=min_samples)
neigh.fit(matrix_scaled)
distances, _ = neigh.kneighbors(matrix_scaled)
kth_distances = distances[:, -1]
eps = float(np.percentile(kth_distances, 90)) # choose a high‑ish value.
print(f"DBSCAN min_samples={min_samples}, eps={eps:.3f}", flush=True)
model = DBSCAN(eps=eps, min_samples=min_samples)
return model.fit_predict(matrix_scaled)
# ---------------------------------------------------------------------------
# Cluster labelling helpers (LLM)
# ---------------------------------------------------------------------------
def label_clusters(
df: pd.DataFrame, labels: np.ndarray, chat_model: str, max_examples: int = 12
) -> dict[int, dict[str, str]]:
"""Generate a name & description for each cluster label via ChatGPT.
Returns a mapping ``label -> {"name": str, "description": str}``.
"""
openai = _lazy_import_openai()
client = openai.OpenAI()
out: dict[int, dict[str, str]] = {}
for lbl in sorted(set(labels)):
if lbl == -1:
# Noise (DBSCAN) – skip LLM call.
out[lbl] = {
"name": "Noise / Outlier",
"description": "Prompts that do not cleanly belong to any cluster.",
}
continue
# Pick a handful of example prompts to send to the model.
examples_series = df.loc[labels == lbl, "prompt"].sample(
min(max_examples, (labels == lbl).sum()), random_state=42
)
examples = examples_series.tolist()
user_content = (
"The following text snippets are all part of the same semantic cluster.\n"
"Please propose \n"
"1. A very short *title* for the cluster (≤ 4 words).\n"
"2. A concise 2–3 sentence *description* that explains the common theme.\n\n"
"Answer **strictly** as valid JSON with the keys 'name' and 'description'.\n\n"
"Snippets:\n"
)
user_content += "\n".join(f"- {t}" for t in examples)
messages = [
{
"role": "system",
"content": "You are an expert analyst, competent in summarising text clusters succinctly.",
},
{"role": "user", "content": user_content},
]
try:
resp = client.chat.completions.create(model=chat_model, messages=messages)
reply = resp.choices[0].message.content.strip()
# Extract the JSON object even if the assistant wrapped it in markdown
# code fences or added other text.
# Remove common markdown fences.
reply_clean = reply.strip()
# Take the substring between the first "{" and the last "}".
m_start = reply_clean.find("{")
m_end = reply_clean.rfind("}")
if m_start == -1 or m_end == -1:
raise ValueError("No JSON object found in model reply.")
json_str = reply_clean[m_start : m_end + 1]
data = json.loads(json_str) # type: ignore[arg-type]
out[lbl] = {
"name": str(data.get("name", "Unnamed"))[:60],
"description": str(data.get("description", "")).strip(),
}
except Exception as exc: # pragma: no cover – network / runtime errors.
print(f"⚠️ Failed to label cluster {lbl}: {exc}", file=sys.stderr)
out[lbl] = {"name": f"Cluster {lbl}", "description": "<LLM call failed>"}
return out
# ---------------------------------------------------------------------------
# Reporting helpers
# ---------------------------------------------------------------------------
def generate_markdown_report(
df: pd.DataFrame,
labels: np.ndarray,
meta: dict[int, dict[str, str]],
outputs: dict[str, Any],
path_md: Path,
):
"""Write a self‑contained Markdown analysis to *path_md*."""
path_md.parent.mkdir(parents=True, exist_ok=True)
cluster_ids = sorted(set(labels))
counts = {lbl: int((labels == lbl).sum()) for lbl in cluster_ids}
lines: list[str] = []
lines.append("# Prompt Clustering Report\n")
lines.append(f"Generated by `cluster_prompts.py` – {pd.Timestamp.now()}\n")
# High‑level stats
total = len(labels)
num_clusters = len(cluster_ids) - (1 if -1 in cluster_ids else 0)
lines.append("\n## Overview\n")
lines.append(f"* Total prompts: **{total}**")
lines.append(f"* Clustering method: **{outputs['method']}**")
if outputs.get("k"):
lines.append(f"* k (K‑Means): **{outputs['k']}**")
lines.append(f"* Silhouette score: **{outputs['silhouette']:.3f}**")
lines.append(f"* Final clusters (excluding noise): **{num_clusters}**\n")
# Summary table
lines.append("\n| label | name | #prompts | description |")
lines.append("|-------|------|---------:|-------------|")
for lbl in cluster_ids:
meta_lbl = meta[lbl]
lines.append(f"| {lbl} | {meta_lbl['name']} | {counts[lbl]} | {meta_lbl['description']} |")
# Detailed section per cluster
for lbl in cluster_ids:
lines.append("\n---\n")
meta_lbl = meta[lbl]
lines.append(f"### Cluster {lbl}: {meta_lbl['name']} ({counts[lbl]} prompts)\n")
lines.append(f"{meta_lbl['description']}\n")
# Show a handful of illustrative prompts.
sample_n = min(5, counts[lbl])
examples = df.loc[labels == lbl, "prompt"].sample(sample_n, random_state=42).tolist()
lines.append("\nExamples:\n")
lines.extend([f"* {t}" for t in examples])
# Outliers / ambiguous prompts, if any.
if -1 in cluster_ids:
lines.append("\n---\n")
lines.append(f"### Noise / outliers ({counts[-1]} prompts)\n")
examples = (
df.loc[labels == -1, "prompt"].sample(min(10, counts[-1]), random_state=42).tolist()
)
lines.extend([f"* {t}" for t in examples])
# Optional ambiguous set (for kmeans)
ambiguous = outputs.get("ambiguous", [])
if ambiguous:
lines.append("\n---\n")
lines.append(f"### Potentially ambiguous prompts ({len(ambiguous)})\n")
lines.extend([f"* {t}" for t in ambiguous])
# Plot references
lines.append("\n---\n")
lines.append("## Plots\n")
lines.append(
"The directory `plots/` contains a bar chart of the cluster sizes and a t‑SNE scatter plot coloured by cluster.\n"
)
path_md.write_text("\n".join(lines))
# ---------------------------------------------------------------------------
# Plotting helpers
# ---------------------------------------------------------------------------
def create_plots(
matrix: np.ndarray,
labels: np.ndarray,
for_devs: pd.Series | None,
plots_dir: Path,
):
"""Generate cluster size and t‑SNE plots."""
import matplotlib.pyplot as plt # type: ignore – heavy, lazy import.
from sklearn.manifold import TSNE # type: ignore – heavy, lazy import.
plots_dir.mkdir(parents=True, exist_ok=True)
# Bar chart with cluster sizes
unique, counts = np.unique(labels, return_counts=True)
order = np.argsort(-counts) # descending
unique, counts = unique[order], counts[order]
plt.figure(figsize=(8, 4))
plt.bar([str(u) for u in unique], counts, color="steelblue")
plt.xlabel("Cluster label")
plt.ylabel("# prompts")
plt.title("Cluster sizes")
plt.tight_layout()
bar_path = plots_dir / "cluster_sizes.png"
plt.savefig(bar_path, dpi=150)
plt.close()
# t‑SNE scatter
tsne = TSNE(
n_components=2, perplexity=min(30, len(matrix) // 3), random_state=42, init="random"
)
xy = tsne.fit_transform(matrix)
plt.figure(figsize=(7, 6))
scatter = plt.scatter(xy[:, 0], xy[:, 1], c=labels, cmap="tab20", s=20, alpha=0.8)
plt.title("t‑SNE projection")
plt.xticks([])
plt.yticks([])
if for_devs is not None:
# Overlay dev prompts as black edge markers
dev_mask = for_devs.astype(bool).values
plt.scatter(
xy[dev_mask, 0],
xy[dev_mask, 1],
facecolors="none",
edgecolors="black",
linewidths=0.6,
s=40,
label="for_devs = TRUE",
)
plt.legend(loc="best")
tsne_path = plots_dir / "tsne.png"
plt.tight_layout()
plt.savefig(tsne_path, dpi=150)
plt.close()
# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
def main() -> None: # noqa: D401
args = parse_cli()
# Read CSV – require a 'prompt' column.
df = pd.read_csv(args.csv)
if "prompt" not in df.columns:
raise SystemExit("Input CSV must contain a 'prompt' column.")
# Keep relevant columns only for clarity.
df = df[[c for c in df.columns if c in {"act", "prompt", "for_devs"}]]
# ---------------------------------------------------------------------
# 1. Embeddings (may be cached)
# ---------------------------------------------------------------------
embeddings_df = load_or_create_embeddings(
df["prompt"], cache_path=args.cache, model=args.embedding_model
)
# ---------------------------------------------------------------------
# 2. Clustering
# ---------------------------------------------------------------------
mat = embeddings_df.values.astype(np.float32)
if args.cluster_method == "kmeans":
labels = cluster_kmeans(mat, k_max=args.k_max)
else:
labels = cluster_dbscan(mat, min_samples=args.dbscan_min_samples)
# Identify potentially ambiguous prompts (only meaningful for kmeans).
outputs: dict[str, Any] = {"method": args.cluster_method}
if args.cluster_method == "kmeans":
from sklearn.cluster import KMeans # type: ignore – lazy
best_k = len(set(labels))
# Re‑fit KMeans with the chosen k to get distances.
kmeans = KMeans(n_clusters=best_k, random_state=42, n_init="auto").fit(mat)
outputs["k"] = best_k
# Silhouette score (again) – not super efficient but okay.
from sklearn.metrics import silhouette_score # type: ignore
outputs["silhouette"] = silhouette_score(mat, labels)
distances = kmeans.transform(mat)
# Ambiguous if the ratio between 1st and 2nd closest centroid < 1.1
sorted_dist = np.sort(distances, axis=1)
ratio = sorted_dist[:, 0] / (sorted_dist[:, 1] + 1e-9)
ambiguous_mask = ratio > 0.9 # tunes threshold – close centroids.
outputs["ambiguous"] = df.loc[ambiguous_mask, "prompt"].tolist()
# ---------------------------------------------------------------------
# 3. LLM naming / description
# ---------------------------------------------------------------------
meta = label_clusters(df, labels, chat_model=args.chat_model)
# ---------------------------------------------------------------------
# 4. Plots
# ---------------------------------------------------------------------
create_plots(mat, labels, df.get("for_devs"), args.plots_dir)
# ---------------------------------------------------------------------
# 5. Markdown report
# ---------------------------------------------------------------------
generate_markdown_report(df, labels, meta, outputs, path_md=args.output_md)
print(f"✅ Done. Report written to {args.output_md} – plots in {args.plots_dir}/", flush=True)
if __name__ == "__main__":
# Guard the main block to allow safe import elsewhere.
main()