pip install ms-swift -U
pip install torch transformers datasets accelerate

from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import mine_hard_negatives


# 加載原始數(shù)據(jù)
dataset = load_dataset("parquet", data_files="/path/to/Chinese-QA-AFAF-train-v2.parquet")
split_dataset = dataset["train"].train_test_split(test_size=0.95, seed=42)  # 劃分訓(xùn)練集和測試集


# 加載基礎(chǔ)模型用于初步檢索
embedding_model = SentenceTransformer("/path/to/m3e-small")


# 挖掘難負(fù)樣本
hard_train_dataset = mine_hard_negatives(
    split_dataset['train'],
    embedding_model,
    anchor_column_name="prompt",  # 查詢列
    positive_column_name="response",  # 正例列
    num_negatives=5,  # 每個(gè)查詢配5個(gè)負(fù)例
    range_min=20,  # 跳過最相似的前20個(gè)（避免與正例太近）
    range_max=50,  # 考慮前50個(gè)相似樣本
    max_score=0.8,  # 負(fù)例與查詢的相似度不超過0.8
    absolute_margin=0.1,  # 負(fù)例相似度需比正例低至少0.1
    use_faiss=True  # 用FAISS加速檢索
)


# 轉(zhuǎn)換為訓(xùn)練格式（查詢+正例+負(fù)例）
def format_data(example):
    correct_resp = next(resp for resp, label in zip(example['response'], example['labels']) if label == 1)
    rejected_resps = [resp for resp, label in zip(example['response'], example['labels']) if label == 0]
    return {
        "query": example['prompt'],
        "response": correct_resp,
        "rejected_response": rejected_resps
    }


# 保存為JSON
hard_train_dataset.map(format_data).to_json("/path/to/qwen3_emb_train.json", force_ascii=False)

{
    "query": "綠色對老年人有哪些健康益處？",
    "response": "綠色環(huán)境能幫助老年人放松神經(jīng)、穩(wěn)定血壓...",
    "rejected_response": [
        "綠茶可預(yù)防心血管疾病...",  // 同含“綠色”但主題不同
        "生態(tài)養(yǎng)殖產(chǎn)出綠色食品..."
    ]
}

# 設(shè)置環(huán)境變量，過濾無效負(fù)例
INFONCE_MASK_FAKE_NEGATIVE=true


# 啟動(dòng)微調(diào)
swift sft \
  --model /path/to/Qwen3-Embedding-0.6B \  # 基礎(chǔ)模型路徑
  --task_type embedding \  # 任務(wù)類型
  --model_type qwen3_emb \  # 模型類型
  --train_type lora \  # 采用LoRA微調(diào)
  --dataset /path/to/qwen3_emb_train.json \  # 訓(xùn)練數(shù)據(jù)
  --split_dataset_ratio 0.05 \  # 5%數(shù)據(jù)作為驗(yàn)證集
  --eval_strategy steps \  # 按步驟驗(yàn)證
  --output_dir ./emb_output \  # 輸出目錄
  --eval_steps 100 \  # 每100步驗(yàn)證一次
  --num_train_epochs 1 \  # 訓(xùn)練輪次
  --save_steps 100 \  # 每100步保存一次
  --per_device_train_batch_size 4 \  # 單卡批次大小
  --gradient_accumulation_steps 4 \  # 梯度累積步數(shù)
  --learning_rate 6e-6 \  # 學(xué)習(xí)率
  --loss_type infonce  # 采用InfoNCE損失

{
    "system": "Judge whether the Document meets the requirements based on the Query and the Instruct. Answer only 'yes' or 'no'.",
    "input": ": Retrieve relevant passages for the query\n: 封閉式陽臺(tái)冬季如何解決春蘭春化問題？\n: 連棟溫室大棚降濕方法：通風(fēng)除濕。",
    "output": "\n\n\nno"  // 文檔與查詢無關(guān)
}

swift sft \
  --model /path/to/Qwen3-Reranker-0.6B \  # 基礎(chǔ)模型路徑
  --dataset /path/to/qwen3_rerank_train.json \  # 訓(xùn)練數(shù)據(jù)
  --model_type qwen3 \  # 模型類型
  --split_dataset_ratio 0.05 \  # 驗(yàn)證集比例
  --torch_dtype bfloat16 \  # 數(shù)據(jù)類型，節(jié)省顯存
  --lora_rank 8 \  # LoRA秩，控制參數(shù)規(guī)模
  --lora_alpha 32 \  # LoRA縮放因子
  --target_modules all-linear \  # 微調(diào)所有線性層
  --eval_strategy steps \
  --output_dir ./rerank_output \
  --max_steps 1000 \  # 最大訓(xùn)練步數(shù)
  --eval_steps 100 \
  --num_train_epochs 1 \
  --save_steps 100 \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 4 \
  --learning_rate 6e-6 \  # 學(xué)習(xí)率
  --loss_scale ignore_empty_think \  # 忽略無效思考過程的損失
  --system 'Judge whether the Document meets the requirements. Answer only "yes" or "no".'  # 系統(tǒng)提示，約束輸出格式