ai_classify 每次通話最多可接受 500 張標籤。 對於較大的分類體系，先針對每份文件依嵌入相似度預先篩選標籤，然後對前 K 個候選項的短名單呼叫 ai_classify。 本教學將教你如何找到最佳 K——即能保持準確度的最小候選數。
Note
NEAREST BY 需要 Databricks Runtime 18 版 或更新版本，或使用無伺服器運算。 Databricks Runtime 18 比 Databricks Runtime 18.0、18.1 和 18.2 更新。
開始之前
- 已啟用 Unity Catalog 的工作區，且可存取
ai_classify（請參閱可用性）。 - Databricks Runtime 18+ 或 Serverless（必須用於
NEAREST BY）。 - 包含待分類文件的 Delta 資料表。
- 一個帶有鍵欄位和可選描述欄位的 Delta 標籤表。
- 一個帶有基層真實標籤的小型評估集，或能夠建立一個（步驟 4 中的選項 B）。
0. 配置
設定你的表格名稱、欄位名稱和嵌入模型。 輔助函式 top_k_labels_json 會為每個文件的前 K 個候選項目建立傳遞給 ai_classify 的 JSON 表達式。
# -- Your tables --
DOCS_TABLE = "path.to.your_docs_table" # table of documents to classify
DOCS_TEXT_COL = "document" # column with text to classify
DOCS_ID_COL = None # unique ID column; set to None to auto-generate via md5
LABELS_TABLE = "path.to.your_labels_table" # table of labels
LABELS_KEY_COL = "label" # column with label value
LABELS_DESC_COL = "description" # description column; set to None if labels have no descriptions
# -- Embedding model --
EMBEDDING_MODEL = "databricks-qwen3-embedding-0-6b" # compact model, good default for English text
# -- K values to sweep --
K_VALUES = [10, 20, 50, 100, 200, 500]
# -- Eval set size (if you need to create one) --
EVAL_SAMPLE_SIZE = 100 # docs to sample for manual labeling
doc_id_expr = DOCS_ID_COL if DOCS_ID_COL else f"md5({DOCS_TEXT_COL})"
label_embed_text = (
f"concat({LABELS_KEY_COL}, ': ', {LABELS_DESC_COL})"
if LABELS_DESC_COL
else LABELS_KEY_COL
)
def top_k_labels_json(prefix=""):
"""Build a JSON expression for collected labels from NEAREST BY results."""
col_prefix = f"{prefix}." if prefix else ""
if LABELS_DESC_COL:
return f"to_json(map_from_entries(collect_list(struct({col_prefix}{LABELS_KEY_COL}, {col_prefix}{LABELS_DESC_COL}))))"
else:
return f"to_json(collect_list({col_prefix}{LABELS_KEY_COL}))"
print(f"Docs table: {DOCS_TABLE} (text: {DOCS_TEXT_COL}, id: {doc_id_expr})")
print(f"Labels table: {LABELS_TABLE} (key: {LABELS_KEY_COL}, desc: {LABELS_DESC_COL})")
print(f"Embed text: {label_embed_text}")
print(f"K sweep: {K_VALUES}")
1. 嵌入標籤
將此執行一次。 只有在分類法改變時才重新運行。
spark.sql(f"""
CREATE OR REPLACE TABLE label_embeddings AS
SELECT
{LABELS_KEY_COL},
{f'{LABELS_DESC_COL},' if LABELS_DESC_COL else ''}
cast(
ai_query('{EMBEDDING_MODEL}', {label_embed_text}) AS ARRAY<FLOAT>
) AS embedding
FROM {LABELS_TABLE}
""")
label_count = spark.sql("SELECT count(*) AS n FROM label_embeddings").first()["n"]
print(f"Embedded {label_count} labels")
2. 嵌入文件
spark.sql(f"""
CREATE OR REPLACE TABLE doc_embeddings AS
SELECT
{doc_id_expr} AS id,
{DOCS_TEXT_COL} AS doc_text,
cast(
ai_query('{EMBEDDING_MODEL}', {DOCS_TEXT_COL}) AS ARRAY<FLOAT>
) AS embedding
FROM {DOCS_TABLE}
""")
doc_count = spark.sql("SELECT count(*) AS n FROM doc_embeddings").first()["n"]
print(f"Embedded {doc_count} documents")
3. 使用 NEAREST BY 檢索頂 K 標籤
NEAREST BY 直接執行近似最近鄰連接——無需中間的 N×M 表格。 對每個文件，它一次回傳 K 個最相似的標籤。
# Preview: top-5 nearest labels for a sample of documents
preview_df = spark.sql(f"""
SELECT
d.id,
l.{LABELS_KEY_COL}
{f', l.{LABELS_DESC_COL}' if LABELS_DESC_COL else ''}
FROM doc_embeddings d
INNER JOIN label_embeddings l
APPROX NEAREST 5 BY SIMILARITY vector_cosine_similarity(d.embedding, l.embedding)
LIMIT 20
""")
preview_df.display()
4. 準備一份標準答案評估資料集
K-tuning 需要少量具有已知正確標籤的文件。 如果你已有評估表，請設 EVAL_TABLE 在下一個格子，跳過抽樣格子。 如果沒有，第二個儲存格會抽樣文件，讓您可以手動加上標籤並重新匯入。
# Option A: point to your existing eval table
# Must have columns: id (matching doc_embeddings.id) and ground_truth_label
EVAL_TABLE = dbutils.widgets.get("eval_table") # read from notebook widget
if EVAL_TABLE:
eval_df = spark.table(EVAL_TABLE)
print(f"Loaded {eval_df.count()} eval examples from {EVAL_TABLE}")
else:
print("No eval table set — run the next cell to sample documents for labeling.")
# Option B: sample documents for manual labeling
if not EVAL_TABLE:
sample_df = spark.sql(f"""
SELECT id, doc_text
FROM doc_embeddings
ORDER BY rand()
LIMIT {EVAL_SAMPLE_SIZE}
""")
sample_df.display()
print(f"\nSampled {EVAL_SAMPLE_SIZE} documents.")
print("Next steps:")
print(" 1. Export these rows (copy the table above or save to CSV)")
print(" 2. Add a 'ground_truth_label' column and fill in the correct label for each doc")
print(" 3. Re-import as a Delta table and set EVAL_TABLE above")
print(" 4. Re-run cell 4 (Option A) to load it")
5. 衡量 Recall@K
Recall@K 會檢查真實標籤是否出現在前 K 個嵌入候選項中。 這是一個 僅供檢索的指標 ——它不會呼叫 ai_classify ，而是立即執行。
如果在特定的 K 值下召回率偏低，ai_classify 就不可能回傳正確答案，因為正確標籤在分類尚未開始前就已被排除在候選集合之外。
assert EVAL_TABLE, "Set EVAL_TABLE in cell 4 before running K-tuning."
spark.sql(f"CREATE OR REPLACE TEMP VIEW eval_set AS SELECT * FROM {EVAL_TABLE}")
recall_results = []
for k in K_VALUES:
row = spark.sql(f"""
SELECT
{k} AS k,
count(*) AS eval_size,
sum(CASE WHEN hit THEN 1 ELSE 0 END) AS hits,
round(sum(CASE WHEN hit THEN 1 ELSE 0 END) / count(*), 4) AS recall_at_k
FROM (
SELECT
e.id,
array_contains(
collect_list(l.{LABELS_KEY_COL}),
e.ground_truth_label
) AS hit
FROM eval_set e
JOIN doc_embeddings d ON d.id = e.id
INNER JOIN label_embeddings l
APPROX NEAREST {k} BY SIMILARITY vector_cosine_similarity(d.embedding, l.embedding)
GROUP BY e.id, e.ground_truth_label
)
""").first()
recall_results.append(row.asDict())
print(f" K={k:>4d} → Recall@K = {row['recall_at_k']:.2%} ({row['hits']}/{row['eval_size']})")
recall_df = spark.createDataFrame(recall_results)
recall_df.display()
6. 衡量端到端的準確度
對每個 K 值，依評估文件建立頂尖 K 標籤集，執行 ai_classify，並與實際情況比較。
此步驟會呼叫 ai_classify，而且成本比召回檢查更高。 從召回率已經相當合理的 K 值開始。
accuracy_results = []
for k in K_VALUES:
# Get top-K labels per eval doc using NEAREST BY
spark.sql(f"""
CREATE OR REPLACE TEMP VIEW eval_top_labels AS
SELECT
d.id,
{top_k_labels_json('l')} AS labels
FROM eval_set e
JOIN doc_embeddings d ON d.id = e.id
INNER JOIN label_embeddings l
APPROX NEAREST {k} BY SIMILARITY vector_cosine_similarity(d.embedding, l.embedding)
GROUP BY d.id
""")
# Materialize ai_classify first (returns VARIANT, and is non-deterministic so can't go inside aggregate)
spark.sql(f"""
CREATE OR REPLACE TEMP VIEW eval_predictions AS
SELECT
e.id,
e.ground_truth_label,
get_json_object(cast(ai_classify(d.doc_text, t.labels, map('version', '2.0')) as string), '$.response[0]') AS predicted_label
FROM eval_set e
JOIN doc_embeddings d ON d.id = e.id
JOIN eval_top_labels t ON t.id = e.id
""")
row = spark.sql(f"""
SELECT
{k} AS k,
count(*) AS eval_size,
sum(CASE WHEN predicted_label = ground_truth_label THEN 1 ELSE 0 END) AS correct,
round(
sum(CASE WHEN predicted_label = ground_truth_label THEN 1 ELSE 0 END) / count(*),
4
) AS accuracy
FROM eval_predictions
""").first()
accuracy_results.append(row.asDict())
print(f" K={k:>4d} → Accuracy = {row['accuracy']:.2%} ({row['correct']}/{row['eval_size']})")
accuracy_df = spark.createDataFrame(accuracy_results)
accuracy_df.display()
7. 比較結果並選擇 K
下圖並列顯示 Recall@K 與端到端準確率。 選擇準確度不再提升時的最小 K 值——K 越大，分類速度越慢，品質卻沒有提升。
import pandas as pd
import matplotlib.pyplot as plt
recall_pd = pd.DataFrame(recall_results)
accuracy_pd = pd.DataFrame(accuracy_results)
combined = recall_pd.merge(accuracy_pd, on="k", suffixes=("_recall", "_acc"))
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(combined["k"], combined["recall_at_k"], "o-", label="Recall@K", linewidth=2)
ax.plot(combined["k"], combined["accuracy"], "s--", label="End-to-end accuracy", linewidth=2)
ax.set_xlabel("K (candidate labels per document)")
ax.set_ylabel("Score")
ax.set_title("K-Tuning: Recall@K vs End-to-End Accuracy")
ax.set_ylim(0, 1.05)
ax.set_xticks(combined["k"])
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
print("\nFull results:")
print(combined[["k", "recall_at_k", "accuracy"]].to_string(index=False))
# Pick your K based on the chart above
CHOSEN_K = 50 # <-- edit this
chosen_row = combined[combined["k"] == CHOSEN_K].iloc[0]
print(f"Chosen K = {CHOSEN_K}")
print(f" Recall@K: {chosen_row['recall_at_k']:.2%}")
print(f" End-to-end accuracy: {chosen_row['accuracy']:.2%}")
8. 對所選 K 進行完整分類
將選取的 K 套用到整個文件表。
spark.sql(f"""
CREATE TABLE IF NOT EXISTS top_labels_per_doc AS
SELECT
d.id,
{top_k_labels_json('l')} AS labels
FROM doc_embeddings d
INNER JOIN label_embeddings l
APPROX NEAREST {CHOSEN_K} BY SIMILARITY vector_cosine_similarity(d.embedding, l.embedding)
GROUP BY d.id
""")
print(f"Built top-{CHOSEN_K} label sets for all documents")
result_df = spark.sql(f"""
SELECT
c.{DOCS_TEXT_COL},
cast(ai_classify(c.{DOCS_TEXT_COL}, t.labels, map('version', '2.0')) as string) AS classification
FROM {DOCS_TABLE} c
JOIN top_labels_per_doc t ON t.id = {doc_id_expr.replace(DOCS_TEXT_COL, f'c.{DOCS_TEXT_COL}')}
""")
result_df.display()
# Optionally save results
# result_df.write.mode("overwrite").saveAsTable("my_catalog.my_schema.classification_results")