hf_train.py

"""
HuggingFace 模型微调训练脚本
============================
使用方法：
  1. 打开 https://colab.research.google.com
  2. Runtime -> Change runtime type -> T4 GPU
  3. 新建一个 Cell，粘贴全部代码运行
  4. 或保存为 .py 文件后 !python train.py 运行

支持场景：classification / sft / dpo
修改 SCENE 变量切换场景
"""

# ============================================================
# 第一步：安装依赖
# ============================================================
import subprocess, sys

def pip_install(packages):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q"] + packages.split())

# 基础依赖（必须）
pip_install("transformers datasets accelerate bitsandbytes sentencepiece scikit-learn")

# 场景依赖（按需取消注释）
# pip_install("trl peft")          # SFT 和 DPO 需要
# pip_install("wandb")             # 训练可视化（可选）

print("✅ 依赖安装完成")

# ============================================================
# 第二步：登录 HuggingFace
# ============================================================
HF_TOKEN = "YOUR_HF_TOKEN_HERE"   # <-- 替换成你的 token

from huggingface_hub import login
login(token=HF_TOKEN)
print("✅ HuggingFace 登录成功")

# ============================================================
# 第三步：选择场景
# ============================================================
SCENE = "classification"   # "classification" | "sft" | "dpo"
print(f"✅ 选择场景：{SCENE}")

# ============================================================
# 场景1：文本分类（金融情感分析）
# ============================================================
if SCENE == "classification":
    import torch
    import numpy as np
    from transformers import (
        AutoTokenizer, AutoModelForSequenceClassification,
        TrainingArguments, Trainer
    )
    from datasets import load_dataset
    from sklearn.metrics import accuracy_score, f1_score

    # --- 配置 ---
    MODEL_NAME = "bert-base-chinese"              # 中文基础模型
    NUM_LABELS = 3                                 # 3分类：正/负/中
    OUTPUT_DIR = "./fin-sentiment-model"
    HF_REPO = "Dshane26/fin-sentiment-cn"          # 上传目标仓库

    # --- 加载数据集 ---
    # 使用 HF 上的金融情感数据集（英文）
    # 中文数据请替换为自己的 CSV：
    #   dataset = load_dataset("csv", data_files="my_data.csv", split="train")
    dataset = load_dataset("financial_phrasebank", "sentences_allagree", split="train")

    print(f"数据集大小：{len(dataset)}")
    print(f"字段：{dataset.column_names}")
    print(f"样例：{dataset[0]}")

    # --- 加载模型 ---
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=NUM_LABELS)

    # --- 数据预处理 ---
    def preprocess(examples):
        return tokenizer(examples["sentence"], truncation=True, padding="max_length", max_length=128)

    dataset = dataset.map(preprocess, batched=True)
    dataset = dataset.rename_column("label", "labels")
    dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

    split = dataset.train_test_split(test_size=0.15, seed=42)
    train_ds, val_ds = split["train"], split["test"]
    print(f"训练集：{len(train_ds)} | 验证集：{len(val_ds)}")

    # --- 评估指标 ---
    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        preds = np.argmax(logits, axis=-1)
        return {"accuracy": accuracy_score(labels, preds), "f1_macro": f1_score(labels, preds, average="macro")}

    # --- 训练参数 ---
    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=3,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=32,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_ratio=0.1,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1_macro",
        fp16=True,
        report_to="none",
        logging_steps=50,
        push_to_hub=True,
        hub_model_id=HF_REPO,
        hub_strategy="end",
    )

    # --- 开始训练 ---
    trainer = Trainer(
        model=model, args=training_args,
        train_dataset=train_ds, eval_dataset=val_ds,
        compute_metrics=compute_metrics,
    )
    trainer.train()
    print("\n✅ 训练完成！")

    # --- 评估 ---
    results = trainer.evaluate()
    print("\n📊 最终评估结果：")
    for k, v in results.items():
        print(f"  {k}: {v:.4f}")

    # --- 推理测试 ---
    label_map = {0: "负面 😞", 1: "中性 😐", 2: "正面 😊"}
    test_texts = ["公司营收增长超预期，利润大幅提升", "业绩持续下滑，亏损扩大", "公司季度报告如预期发布"]
    print("\n🧪 推理测试：")
    for text in test_texts:
        inputs = tokenizer(text, return_tensors="pt", truncation=True).to(model.device)
        with torch.no_grad():
            pred = torch.argmax(model(**inputs).logits, dim=-1).item()
        print(f"  「{text}」→ {label_map[pred]}")

    print(f"\n🎉 模型已上传到 https://huggingface.co/{HF_REPO}")


# ============================================================
# 场景2：指令微调 SFT
# ============================================================
elif SCENE == "sft":
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, BitsAndBytesConfig
    from trl import SFTTrainer
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from datasets import load_dataset

    BASE_MODEL = "Qwen/Qwen2-1.5B-Instruct"
    OUTPUT_DIR = "./qwen-sft-model"
    HF_REPO = "Dshane26/qwen-sft-finance"

    # 4-bit 量化加载
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL, quantization_config=bnb_config, device_map="auto", trust_remote_code=True
    )
    model = prepare_model_for_kbit_training(model)
    model = get_peft_model(model, LoraConfig(
        r=16, lora_alpha=32,
        target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
        lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    ))
    model.print_trainable_parameters()

    # 加载数据
    dataset = load_dataset("tatsu-lab/alpaca", split="train")
    def format_alpaca(ex):
        if ex["input"]:
            text = f"### Instruction:\n{ex['instruction']}\n\n### Input:\n{ex['input']}\n\n### Response:\n{ex['output']}"
        else:
            text = f"### Instruction:\n{ex['instruction']}\n\n### Response:\n{ex['output']}"
        return {"text": text}
    dataset = dataset.map(format_alpaca)
    print(f"数据集大小：{len(dataset)}")

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR, num_train_epochs=3,
        per_device_train_batch_size=4, gradient_accumulation_steps=4,
        learning_rate=2e-4, weight_decay=0.01, warmup_ratio=0.05,
        lr_scheduler_type="cosine", logging_steps=100,
        save_strategy="epoch", fp16=True, report_to="none",
        push_to_hub=True, hub_model_id=HF_REPO, hub_strategy="end",
    )

    trainer = SFTTrainer(
        model=model, args=training_args, train_dataset=dataset,
        dataset_text_field="text", max_seq_length=512, tokenizer=tokenizer,
    )
    trainer.train()
    print("\n✅ 训练完成！")

    # 推理测试
    prompt = "### Instruction:\n解释什么是市盈率（PE Ratio）\n\n### Response:\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.7, do_sample=True)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"\n🧪 推理测试：\n{result.split('### Response:')[-1].strip()}")
    print(f"\n🎉 模型已上传到 https://huggingface.co/{HF_REPO}")


# ============================================================
# 场景3：DPO 对齐
# ============================================================
elif SCENE == "dpo":
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, BitsAndBytesConfig
    from trl import DPOTrainer
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from datasets import load_dataset

    BASE_MODEL = "Qwen/Qwen2-1.5B-Instruct"
    OUTPUT_DIR = "./qwen-dpo-model"
    HF_REPO = "Dshane26/qwen-dpo-finance"

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL, quantization_config=bnb_config, device_map="auto", trust_remote_code=True
    )
    model_ref = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL, quantization_config=bnb_config, device_map="auto", trust_remote_code=True
    )
    model = prepare_model_for_kbit_training(model)
    model = get_peft_model(model, LoraConfig(
        r=16, lora_alpha=32, target_modules=["q_proj","k_proj","v_proj","o_proj"],
        lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    ))

    dataset = load_dataset("tatsu-lab/alpaca", split="train")
    print(f"数据集大小：{len(dataset)}")

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR, num_train_epochs=1,
        per_device_train_batch_size=2, gradient_accumulation_steps=8,
        learning_rate=5e-7, beta=0.1, logging_steps=50,
        save_strategy="epoch", fp16=True, report_to="none",
        push_to_hub=True, hub_model_id=HF_REPO, hub_strategy="end",
    )

    trainer = DPOTrainer(
        model=model, ref_model=model_ref, args=training_args,
        train_dataset=dataset, tokenizer=tokenizer, max_length=512, beta=0.1,
    )
    trainer.train()
    print("\n✅ DPO 训练完成！")
    print(f"\n🎉 模型已上传到 https://huggingface.co/{HF_REPO}")

else:
    print(f"❌ 未知场景：{SCENE}。请选择 classification / sft / dpo")