Embedding 模型微调实战（m3e-base）

基于 moka-ai/m3e-base 中文 Embedding 模型的本地微调 demo。

环境要求

• Python 3.9+
• 内存 >= 8GB（推荐 16GB）
• 无需 GPU，CPU 即可运行

快速开始

# 1. 创建项目目录
mkdir embedding-finetune
cd embedding-finetune

# 2. 创建虚拟环境
python3 -m venv venv
source venv/bin/activate        # macOS / Linux
# venv\Scripts\activate         # Windows

# 3. 创建 requirements.txt
cat > requirements.txt << 'EOF'
sentence-transformers>=3.0.0
torch>=2.0.0
datasets>=2.14.0
scikit-learn>=1.3.0
pandas>=2.0.0
EOF

# 4. 安装依赖（装在 venv 里，不影响全局环境）
pip install -r requirements.txt

# 5. 如果 HuggingFace 下载慢，设置国内镜像
export HF_ENDPOINT=https://hf-mirror.com

然后将 step1~step4 脚本和 train_data.csv 放入项目目录即可。

操作步骤

按顺序执行 4 个脚本：

Step 1: 测试模型能否正常运行

python step1_test_model.py

首次运行会自动下载 m3e-base 模型（约 400MB）。
验证模型加载和推理是否正常。

Step 2: 查看训练数据

python step2_prepare_data.py

读取 train_data.csv，展示数据统计和分布。
你可以修改 train_data.csv 添加自己的业务数据。

Step 3: 微调模型

python step3_finetune.py

CPU 上约 5~10 分钟完成。
微调后的模型保存在 output/finetuned-m3e/final/。

Step 4: 对比效果

python step4_compare.py

对比原始模型和微调模型在同一组测试数据上的相似度差异。

训练数据格式

train_data.csv 格式：

sentence1,sentence2,score
杭州西湖大酒店,西湖大饭店杭州,0.95
杭州西湖大酒店,上海浦东希尔顿,0.15
标准间,标间,1.0

• sentence1, sentence2: 两条文本
• score: 0~1 的相似度分数（1=完全相同，0=完全无关）

数据准备建议

• 至少准备 100 条以上，效果更好
• 正样本（高相似度）和负样本（低相似度）都要有
• 多放"容易混淆"的 case，这是模型最需要学的
• 数据质量 > 数量

项目结构

embedding-finetune/
├── README.md                  # 本文件
├── requirements.txt           # Python 依赖
├── train_data.csv             # 训练数据（可自行修改）
├── step1_test_model.py        # 测试模型推理
├── step2_prepare_data.py      # 查看训练数据
├── step3_finetune.py          # 微调训练
├── step4_compare.py           # 效果对比
└── output/                    # 训练输出（自动生成）
    └── finetuned-m3e/
        └── final/             # 最终模型

常见问题

下载模型很慢？

# 使用 HuggingFace 国内镜像
export HF_ENDPOINT=https://hf-mirror.com

内存不够（OOM）？

修改 step3_finetune.py 中的 per_device_train_batch_size，从 8 改为 4 或 2。

Mac M 系列芯片加速？

取消 step3_finetune.py 中 use_mps_device=True 的注释。

数据准备

sentence1,sentence2,score
杭州西湖大酒店,西湖大饭店杭州,0.95
杭州西湖大酒店,杭州西湖大饭店,0.98
杭州西湖大酒店,上海浦东希尔顿酒店,0.15
北山街38号,北山路38号,0.95
杭州市西湖区,浙江省杭州市西湖区,0.92
杭州市西湖区北山街,西湖区北山路杭州,0.90
标准间,标间,1.0
大床房,豪华大床房,0.85
双人间,双人标准间,0.90
行政套房,商务套房,0.80
含早餐,含双早,0.90
含早餐,不含餐,0.10
五星级酒店,豪华酒店,0.80
经济型酒店,快捷酒店,0.90
商务酒店,商务型宾馆,0.90
酒店,饭店,0.85
酒店,宾馆,0.90
酒店,旅馆,0.80
酒店,超市,0.05
酒店前台,前台接待,0.90
房间干净整洁,客房非常整洁,0.92
交通便利,出行方便,0.88
靠近地铁站,地铁站附近,0.95
免费停车,提供免费停车场,0.92
免费WiFi,无线网络免费,0.90
电话0571-87991234,Tel: 0571-8799-1234,0.98
13800138000,138-0013-8000,0.98
杭州萧山国际机场附近酒店,萧山机场旁边的宾馆,0.90
西湖景区旁边酒店,西湖风景区附近饭店,0.92
杭州火车东站酒店,杭州东站附近宾馆,0.90
价格实惠,性价比高,0.80
价格实惠,豪华昂贵,0.10
可携带宠物,宠物友好,0.90
提供接机服务,机场接送,0.88
有游泳池,含泳池设施,0.90
会议室,商务会议厅,0.85
自助早餐,早餐自助餐,0.95
中式早餐,中式早点,0.95
杭州西溪湿地酒店,西溪湿地旁宾馆,0.90
千岛湖度假酒店,千岛湖度假村饭店,0.88
浙江省杭州市拱墅区,杭州拱墅区,0.92
上城区解放路100号,杭州市上城区解放路100号,0.95
河坊街美食酒店,杭州河坊街附近饭店,0.85
南宋御街文化主题酒店,御街文化主题宾馆,0.88

Step 1: 下载 m3e-base 模型并测试推理

"""
Step 1: 下载 m3e-base 模型并测试推理
运行: python step1_test_model.py

首次运行会自动从 HuggingFace 下载模型（约 400MB），之后直接读本地缓存。
如果下载慢，可以设置镜像：
  export HF_ENDPOINT=https://hf-mirror.com
"""

from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity


def main():
    print("=" * 60)
    print("Step 1: 加载 m3e-base 模型")
    print("=" * 60)

    # 加载模型（首次会下载，之后读缓存）
    print("\n正在加载模型（首次需要下载约 400MB）...")
    model = SentenceTransformer("moka-ai/m3e-base")
    print("模型加载成功！\n")

    # 测试单个文本向量化
    text = "杭州西湖大酒店"
    vector = model.encode(text)
    print(f"文本: '{text}'")
    print(f"向量维度: {len(vector)}")
    print(f"前5个值: {vector[:5]}")
    print()

    # 测试多组相似度
    print("-" * 60)
    print("相似度测试")
    print("-" * 60)

    test_pairs = [
        # 语义相近的对（期望高相似度）
        ("杭州西湖大酒店", "西湖大饭店（杭州）"),
        ("北山街38号", "北山路38号"),
        ("标准间", "标间"),
        ("含早餐", "含双早"),
        ("五星级酒店", "豪华酒店"),
        # 语义不相关的对（期望低相似度）
        ("杭州西湖大酒店", "北京烤鸭"),
        ("标准间", "火车票"),
        ("酒店前台", "股票基金"),
    ]

    print(f"\n{'文本A':<16} {'文本B':<16} {'相似度':>8}")
    print("-" * 50)

    for text_a, text_b in test_pairs:
        vec_a = model.encode(text_a)
        vec_b = model.encode(text_b)
        score = cosine_similarity([vec_a], [vec_b])[0][0]
        print(f"{text_a:<16} {text_b:<16} {score:>8.4f}")

    print("\n✅ Step 1 完成！模型可以正常运行。")
    print("👉 下一步: python step2_prepare_data.py")


if __name__ == "__main__":
    main()

Step 2: 查看和验证训练数据

"""
Step 2: 查看和验证训练数据
运行: python step2_prepare_data.py

这个脚本会：
1. 读取 train_data.csv 并展示数据统计
2. 验证数据格式是否正确
3. 展示数据分布情况

你可以根据自己的业务修改 train_data.csv，格式为：
  sentence1,sentence2,score
  - sentence1, sentence2: 两条文本
  - score: 0~1 之间的相似度分数（1=完全相同，0=完全无关）
"""

import csv
import os


def main():
    print("=" * 60)
    print("Step 2: 查看训练数据")
    print("=" * 60)

    data_file = os.path.join(os.path.dirname(__file__), "train_data.csv")

    if not os.path.exists(data_file):
        print(f"\n❌ 找不到 {data_file}")
        print("请先创建 train_data.csv 文件")
        return

    # 读取数据
    rows = []
    with open(data_file, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            rows.append(
                {
                    "sentence1": row["sentence1"],
                    "sentence2": row["sentence2"],
                    "score": float(row["score"]),
                }
            )

    print(f"\n总数据量: {len(rows)} 条")

    # 数据分布
    high = [r for r in rows if r["score"] >= 0.8]
    medium = [r for r in rows if 0.4 <= r["score"] < 0.8]
    low = [r for r in rows if r["score"] < 0.4]

    print(f"\n分布情况:")
    print(f"  高相似度 (>=0.8): {len(high)} 条")
    print(f"  中相似度 (0.4~0.8): {len(medium)} 条")
    print(f"  低相似度 (<0.4):  {len(low)} 条")

    # 展示部分样本
    print(f"\n--- 高相似度样本（前 5 条）---")
    for r in high[:5]:
        print(f"  [{r['score']:.2f}] {r['sentence1']}  ↔  {r['sentence2']}")

    print(f"\n--- 低相似度样本（前 5 条）---")
    for r in low[:5]:
        print(f"  [{r['score']:.2f}] {r['sentence1']}  ↔  {r['sentence2']}")

    # 验证
    errors = []
    for i, r in enumerate(rows):
        if not r["sentence1"].strip():
            errors.append(f"第 {i+2} 行: sentence1 为空")
        if not r["sentence2"].strip():
            errors.append(f"第 {i+2} 行: sentence2 为空")
        if not (0 <= r["score"] <= 1):
            errors.append(f"第 {i+2} 行: score={r['score']} 不在 0~1 范围")

    if errors:
        print(f"\n❌ 发现 {len(errors)} 个问题:")
        for e in errors:
            print(f"  - {e}")
    else:
        print(f"\n✅ 数据格式验证通过！")

    # 建议
    print(f"\n💡 建议:")
    if len(rows) < 50:
        print(f"  - 当前 {len(rows)} 条数据偏少，建议补充到 100+ 条效果更好")
    if len(low) < len(rows) * 0.15:
        print(f"  - 负样本（低相似度）偏少，建议补充一些不相关的文本对")
    if len(rows) >= 50 and len(low) >= len(rows) * 0.15:
        print(f"  - 数据量和分布看起来不错，可以开始微调了！")

    print(f"\n✅ Step 2 完成！")
    print(f"👉 下一步: python step3_finetune.py")


if __name__ == "__main__":
    main()

微调 m3e-base 模型

"""
Step 3: 微调 m3e-base 模型
运行: python step3_finetune.py

这个脚本会：
1. 加载 m3e-base 基础模型
2. 读取 train_data.csv 训练数据
3. 用 CosineSimilarityLoss 进行微调
4. 保存微调后的模型到 ./output/finetuned-m3e/

预计耗时（CPU）：约 5~10 分钟（取决于数据量和机器性能）
预计内存占用：约 3 GB
"""

import csv
import os

from datasets import Dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    losses,
)


def load_train_data(data_file):
    """从 CSV 加载训练数据"""
    rows = []
    with open(data_file, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            rows.append(
                {
                    "sentence1": row["sentence1"],
                    "sentence2": row["sentence2"],
                    "score": float(row["score"]),
                }
            )
    return rows


def main():
    print("=" * 60)
    print("Step 3: 微调 m3e-base 模型")
    print("=" * 60)

    # ========== 1. 加载基础模型 ==========
    print("\n[1/5] 加载 m3e-base 基础模型...")
    model = SentenceTransformer("moka-ai/m3e-base")
    print(f"  模型加载成功，向量维度: {model.get_sentence_embedding_dimension()}")

    # ========== 2. 加载训练数据 ==========
    print("\n[2/5] 加载训练数据...")
    data_file = os.path.join(os.path.dirname(__file__), "train_data.csv")
    rows = load_train_data(data_file)

    dataset = Dataset.from_list(rows)

    # 拆分训练集和验证集（90% / 10%）
    split = dataset.train_test_split(test_size=0.1, seed=42)
    train_dataset = split["train"]
    eval_dataset = split["test"]

    print(f"  训练集: {len(train_dataset)} 条")
    print(f"  验证集: {len(eval_dataset)} 条")

    # ========== 3. 定义损失函数 ==========
    print("\n[3/5] 配置训练参数...")
    # CosineSimilarityLoss: 让模型输出的余弦相似度接近你标注的 score
    loss = losses.CosineSimilarityLoss(model)

    # ========== 4. 训练参数 ==========
    output_dir = os.path.join(os.path.dirname(__file__), "output", "finetuned-m3e")

    args = SentenceTransformerTrainingArguments(
        output_dir=output_dir,
        # --- 训练轮数 ---
        num_train_epochs=10,  # 数据少的时候多训几轮
        # --- 批次大小（内存不够就改小） ---
        per_device_train_batch_size=8,
        # --- 如果 batch_size 改小了，用梯度累积补回来 ---
        gradient_accumulation_steps=2,  # 等效 batch_size = 8 * 2 = 16
        # --- 学习率 ---
        learning_rate=2e-5,
        # --- 预热（前 10% 的步骤逐渐增大学习率，避免一开始太猛） ---
        warmup_ratio=0.1,
        # --- 评估和保存策略 ---
        eval_strategy="epoch",  # 每轮评估一次
        save_strategy="epoch",  # 每轮保存一次
        save_total_limit=2,  # 只保留最近 2 个 checkpoint，节省磁盘
        load_best_model_at_end=True,  # 训练结束后加载最优模型
        # --- 日志 ---
        logging_steps=5,
        # --- Mac M 系列芯片可以用 MPS 加速，取消注释即可 ---
        # use_mps_device=True,
    )

    print(f"  训练轮数: {args.num_train_epochs}")
    print(f"  批次大小: {args.per_device_train_batch_size}")
    print(f"  学习率: {args.learning_rate}")
    print(f"  输出目录: {output_dir}")

    # ========== 5. 开始训练 ==========
    print("\n[4/5] 开始训练...")
    print("  （CPU 上大约 5~10 分钟，请耐心等待）\n")

    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
    )

    trainer.train()

    # ========== 6. 保存最终模型 ==========
    print("\n[5/5] 保存微调后的模型...")
    final_path = os.path.join(output_dir, "final")
    model.save(final_path)
    print(f"  模型已保存到: {final_path}")

    print("\n" + "=" * 60)
    print("✅ 微调完成！")
    print("=" * 60)
    print(f"\n👉 下一步: python step4_compare.py")
    print(f"   对比微调前后的效果差异")


if __name__ == "__main__":
    main()

对比微调前后的效果

"""
Step 4: 对比微调前后的效果
运行: python step4_compare.py

这个脚本会：
1. 分别加载原始 m3e-base 和微调后的模型
2. 用同一组测试数据计算相似度
3. 对比两个模型的差异
"""

import os

from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity


def compute_scores(model, pairs):
    """计算一组文本对的相似度"""
    scores = []
    for text_a, text_b in pairs:
        vec_a = model.encode(text_a)
        vec_b = model.encode(text_b)
        score = cosine_similarity([vec_a], [vec_b])[0][0]
        scores.append(score)
    return scores


def main():
    print("=" * 70)
    print("Step 4: 对比微调前后效果")
    print("=" * 70)

    # 检查微调模型是否存在
    finetuned_path = os.path.join(
        os.path.dirname(__file__), "output", "finetuned-m3e", "final"
    )
    if not os.path.exists(finetuned_path):
        print(f"\n❌ 找不到微调后的模型: {finetuned_path}")
        print("请先运行 python step3_finetune.py")
        return

    # 加载两个模型
    print("\n加载原始模型...")
    original = SentenceTransformer("moka-ai/m3e-base")
    print("加载微调模型...")
    finetuned = SentenceTransformer(finetuned_path)
    print("模型加载完成！\n")

    # 测试数据（包含训练集内和训练集外的样本）
    test_pairs = [
        # --- 训练集中出现过的（验证学习效果）---
        ("标准间", "标间"),
        ("含早餐", "含双早"),
        ("杭州西湖大酒店", "西湖大饭店杭州"),
        ("北山街38号", "北山路38号"),
        # --- 训练集中没出现过的（验证泛化能力）---
        ("海景房", "海景大床房"),
        ("免费接机", "机场免费接送"),
        ("南京东路酒店", "东路饭店南京"),
        ("带浴缸", "含独立浴缸"),
        # --- 不相关的对（应该保持低分）---
        ("酒店入住", "编程语言"),
        ("大床房", "高铁票"),
    ]

    # 计算两个模型的分数
    original_scores = compute_scores(original, test_pairs)
    finetuned_scores = compute_scores(finetuned, test_pairs)

    # 打印对比结果
    print(f"{'文本A':<14} {'文本B':<14} {'原始模型':>8} {'微调模型':>8} {'变化':>8}")
    print("-" * 70)

    for i, (text_a, text_b) in enumerate(test_pairs):
        orig = original_scores[i]
        ft = finetuned_scores[i]
        diff = ft - orig
        arrow = "↑" if diff > 0.01 else ("↓" if diff < -0.01 else "→")
        print(
            f"{text_a:<14} {text_b:<14} {orig:>8.4f} {ft:>8.4f} {arrow}{abs(diff):>6.4f}"
        )

    # 统计
    print("\n" + "=" * 70)
    print("总结")
    print("=" * 70)

    # 相似对（前8个）
    similar_orig = original_scores[:8]
    similar_ft = finetuned_scores[:8]
    avg_orig_sim = sum(similar_orig) / len(similar_orig)
    avg_ft_sim = sum(similar_ft) / len(similar_ft)

    # 不相关对（后2个）
    diff_orig = original_scores[8:]
    diff_ft = finetuned_scores[8:]
    avg_orig_diff = sum(diff_orig) / len(diff_orig)
    avg_ft_diff = sum(diff_ft) / len(diff_ft)

    print(f"\n相似文本对的平均相似度:")
    print(f"  原始模型: {avg_orig_sim:.4f}")
    print(f"  微调模型: {avg_ft_sim:.4f}")

    print(f"\n不相关文本对的平均相似度:")
    print(f"  原始模型: {avg_orig_diff:.4f}")
    print(f"  微调模型: {avg_ft_diff:.4f}")

    print(f"\n理想效果: 相似对分数↑，不相关对分数↓")

    if avg_ft_sim > avg_orig_sim:
        print(f"✅ 相似文本的相似度提升了 {avg_ft_sim - avg_orig_sim:+.4f}")
    else:
        print(f"⚠️  相似文本的相似度下降了 {avg_ft_sim - avg_orig_sim:+.4f}，可能需要调整训练数据或参数")

    print(f"\n💡 提示:")
    print(f"  - 如果效果不明显，尝试增加训练数据量（当前仅为示例数据）")
    print(f"  - 可以在 train_data.csv 中添加更多你业务中的真实数据")
    print(f"  - 然后重新运行 step3 和 step4")


if __name__ == "__main__":
    main()

用 LoRA 微调 m3e-base 模型

代替第三步

"""
Step 3 (LoRA 版): 用 LoRA 微调 m3e-base 模型
运行: python step3_finetune_lora.py

与 step3_finetune.py 的区别：
  - 原版：直接修改模型全部 1.02 亿个参数（全量微调）
  - 本版：冻结原始参数，只训练插入的 LoRA 适配器（约 30 万个参数）

需要额外安装: pip install peft
"""

import csv
import os

from datasets import Dataset
from peft import LoraConfig, TaskType, get_peft_model
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    losses,
)


def load_train_data(data_file):
    """从 CSV 加载训练数据"""
    rows = []
    with open(data_file, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            rows.append(
                {
                    "sentence1": row["sentence1"],
                    "sentence2": row["sentence2"],
                    "score": float(row["score"]),
                }
            )
    return rows


def print_trainable_params(model):
    """打印可训练参数量，直观感受 LoRA 的轻量"""
    transformer = model[0].auto_model
    total = sum(p.numel() for p in transformer.parameters())
    trainable = sum(p.numel() for p in transformer.parameters() if p.requires_grad)
    frozen = total - trainable
    print(f"  总参数量:     {total:>12,}")
    print(f"  冻结参数:     {frozen:>12,}  ({frozen/total*100:.2f}%)")
    print(f"  可训练参数:   {trainable:>12,}  ({trainable/total*100:.2f}%)")


def main():
    print("=" * 60)
    print("Step 3 (LoRA): 用 LoRA 微调 m3e-base 模型")
    print("=" * 60)

    # ========== 1. 加载基础模型 ==========
    print("\n[1/6] 加载 m3e-base 基础模型...")
    model = SentenceTransformer("moka-ai/m3e-base")
    print(f"  向量维度: {model.get_sentence_embedding_dimension()}")

    # ========== 2. 给模型装上 LoRA 适配器 ==========
    print("\n[2/6] 安装 LoRA 适配器...")

    # LoRA 核心配置
    lora_config = LoraConfig(
        task_type=TaskType.FEATURE_EXTRACTION,
        # r: LoRA 的秩（rank），控制适配器的"容量"
        #    r 越大，适配器参数越多，学习能力越强，但也越容易过拟合
        #    通常取 4~16，对于小数据集取小值就够了
        r=8,
        # lora_alpha: 缩放系数，控制 LoRA 的影响力度
        #    通常设为 r 的 2 倍
        lora_alpha=16,
        # lora_dropout: 防止过拟合
        lora_dropout=0.1,
        # target_modules: 在哪些层插入 LoRA
        #    query 和 value 是注意力机制中最关键的两个矩阵
        target_modules=["query", "value"],
    )

    # 把 LoRA 装到模型的 Transformer 层上
    # model[0] 是 SentenceTransformer 的第一个模块（Transformer 编码器）
    model[0].auto_model = get_peft_model(model[0].auto_model, lora_config)

    print_trainable_params(model)

    # ========== 3. 加载训练数据 ==========
    print("\n[3/6] 加载训练数据...")
    data_file = os.path.join(os.path.dirname(__file__), "train_data.csv")
    rows = load_train_data(data_file)

    dataset = Dataset.from_list(rows)
    split = dataset.train_test_split(test_size=0.1, seed=42)
    train_dataset = split["train"]
    eval_dataset = split["test"]

    print(f"  训练集: {len(train_dataset)} 条")
    print(f"  验证集: {len(eval_dataset)} 条")

    # ========== 4. 配置训练 ==========
    print("\n[4/6] 配置训练参数...")
    loss = losses.CosineSimilarityLoss(model)

    output_dir = os.path.join(os.path.dirname(__file__), "output", "finetuned-m3e-lora")

    args = SentenceTransformerTrainingArguments(
        output_dir=output_dir,
        num_train_epochs=10,
        per_device_train_batch_size=8,
        gradient_accumulation_steps=2,
        learning_rate=2e-4,  # LoRA 可以用更大的学习率，因为只动小模块
        warmup_ratio=0.1,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=2,
        load_best_model_at_end=True,
        logging_steps=5,
    )

    # ========== 5. 开始训练 ==========
    print("\n[5/6] 开始训练...")
    print("  （LoRA 比全量微调更快，参数少很多）\n")

    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
    )

    trainer.train()

    # ========== 6. 保存 ==========
    print("\n[6/6] 保存模型...")
    final_path = os.path.join(output_dir, "final")

    # 方式 A：只保存 LoRA 适配器（几 MB）
    lora_path = os.path.join(output_dir, "lora-adapter")
    model[0].auto_model.save_pretrained(lora_path)
    lora_size = sum(
        os.path.getsize(os.path.join(lora_path, f))
        for f in os.listdir(lora_path)
        if os.path.isfile(os.path.join(lora_path, f))
    )
    print(f"  LoRA 适配器已保存到: {lora_path}")
    print(f"  适配器大小: {lora_size / 1024:.1f} KB  (对比全量模型 ~400 MB)")

    # 方式 B：合并成完整模型（方便直接加载，不需要 peft 库）
    model[0].auto_model = model[0].auto_model.merge_and_unload()
    model.save(final_path)
    print(f"  合并后完整模型已保存到: {final_path}")

    print("\n" + "=" * 60)
    print("✅ LoRA 微调完成！")
    print("=" * 60)
    print(f"\n👉 下一步: python step4_compare.py")
    print(f"   (step4 会自动加载 output/finetuned-m3e/final 对比)")
    print(f"   如果想对比 LoRA 版，改 step4 里的路径为:")
    print(f"   output/finetuned-m3e-lora/final")


if __name__ == "__main__":
    main()