Argilla SDK 聊天機器人（使用 distilabel）– 端對端教學

TL;DR

我們使用 distilabel 產生合成的問答三元組，微調帶有 Matryoshka loss 的 BGE‑base 嵌入模型，將嵌入儲存在輕量級 lancedb 向量資料庫，並透過 Gradio 在 Hugging Face Spaces 部署聊天介面，從而為 Argilla 2.0 建立了一個領域特定的 RAG 聊天機器人。

使用 distilabel 產生合成訓練資料

Key outcome: 捕捉 Argilla 文件的真實查詢與困難負樣本的三元組資料集 (anchor, positive, negative).

  • 流程從載入 Hub 資料集 plaguss/argilla_sdk_docs_raw_unstructured 中的原始文件片段開始，並將欄位 chunks 重新命名為 anchor。
  • GenerateSentencePair（三元組模式）使用 LLM meta-llama/Meta-Llama-3-70B-Instruct 為每個片段產生 positive 查詢，並產生一個與之無關但在字面上相似的 negative 查詢。
  • 自訂的 MultipleQueries 任務將每個正向查詢擴展為另外三個變體，使資料集規模增加四倍。
  • MergeColumns 與 ExpandColumns 將原始與產生的查詢合併為單一 positive 欄位，產生每筆 query‑anchor‑negative 三元組對應一列。

完整的流程定義於 pipeline_docs_queries.py，最終資料集會推送至 plaguss/argilla_sdk_docs_queries。

在 Argilla 中探索與整理資料集

Key outcome: 三個 Argilla 資料集，分別用於 (1) 原始文件片段、(2) 嵌入微調三元組、以及 (3) 聊天機器人互動日誌。

  • Documentation chunks – 欄位 filename 與 chunk，以及用於人工驗證的二元標籤 good_chunk。
  • Embedding triples – 欄位 anchor、positive、negative，以及二元相關性標籤 is_positive_relevant 與 is_negative_irrelevant。
  • Chatbot logs – 欄位 instruction 與 response，加上中繼資料 conv_id 與 turn；標籤用於評估正確性與安全欄位違規，且自由格式的 feedback 欄位可捕捉使用者評論。

所有資料集皆透過 Argilla Python 客戶端 (rg.Argilla) 建立，並可直接在 Argilla UI 中檢視。

微調嵌入模型

Key outcome: 一個自訂模型 plaguss/bge-base-argilla-sdk-matryoshka，在 Argilla 專屬檢索上優於基線 BGE‑base。

  1. Dataset preparation – 載入三元組資料集，保留欄位 anchor、positive、negative，新增唯一的 id，並以 90 %/10 % 分割為訓練/測試集。
  2. Baseline model – 以 BAAI/bge-base-en-v1.5 為起點，設定模型卡片的中繼資料。
  3. Loss function – 結合 TripletLoss 與 MatryoshkaLoss（維度 [768, 512, 256, 128, 64]）。
  4. Training arguments – 為 Apple M2 Pro 調整批次大小，使用 cosine 調度器，並選擇 eval_dim_512_cosine_ndcg@10 作為評估指標。
  5. Training – 執行 SentenceTransformerTrainer；最佳檢查點會自動推送至 Hub。

產生的模型可透過 SentenceTransformer 載入，或使用 sentence-transformers 註冊表。

使用 lancedb 建立向量資料庫

Key outcome: 一個可攜帶、無伺服器的向量儲存，將每個合成查詢與其文件片段關聯起來。

  • lancedb.connect("./lancedb") 會建立本機類 SQLite 的資料庫。
  • 使用 LanceModel 定義 Docs 結構（query、text、vector）。
  • 對於查詢資料集的每個批次，使用微調模型產生嵌入，並插入至資料表中。
  • 檢索範例 – 針對 "How can I get the current user?" 進行 cosine 相似度搜尋，返回最相關的文件片段。
  • 整個資料庫目錄會被壓縮為 (lancedb.tar.gz) 並與資料集一起上傳至 Hub，確保可重現下載。

Gradio 聊天介面與部署

Key outcome: 一個互動式網頁 UI (https://huggingface.co/spaces/plaguss/argilla-sdk-chatbot-space)，使用 RAG 回答 Argilla SDK 的問題。

  • Database class – 處理 lancedb 壓縮檔的延遲下載，開啟資料表，並提供 retrieve_doc_chunks，對給定查詢返回最多四個去重的片段。
  • Prompt engineering – 系統提示迫使 LLM 僅從提供的上下文中作答。使用者提示模板（ARGILLA_BOT_TEMPLATE）會插入檢索到的片段。
  • LLM inference – 透過 InferenceClient 呼叫 Hugging Face 推論端點（預設為 meta-llama/Meta-Llama-3-70B-Instruct），回傳的串流結果會傳回給 Gradio。
  • Conversation logging – 每回合結束後，將互動紀錄寫入 Argilla chatbot‑log 資料集，以支援持續評估與未來微調。
  • Deployment – 新增 requirements.txt 並將 Hugging Face API 金鑰設為機密，即可讓應用自動在 Spaces 上建置。

含意與後續步驟

Takeaway: 端對端工作流程展示了如何將任何以程式碼為中心的文件庫轉換為高品質、領域特定的 RAG 聊天機器人，且僅需最少的人工標記。

  • Scalability – 只要更換 GitHub 倉庫路徑，即可將相同流程套用至其他函式庫或內部 SDK。
  • Data quality – 改善片段大小、加入去重、以及強化合成查詢產生（例如使用結構化提示），皆可提升檢索相關性。
  • Explainability – 為返回的片段加入來源 URL 或行號，可提供使用者可追溯性。
  • Feedback loop – Argilla 互動資料集提供即用的回饋迴路，以進行模型的迭代改進。

結合 distilabel 的合成資料生成、Matryoshka 增強的嵌入微調，以及輕量級向量儲存，開發者能快速為任何技術產品原型出可靠的支援機器人。

Sources