# 初始化一個名為 demo-unsloth 的項(xiàng)目，并指定 Python 版本為 3.11.9
uv init demo-unsloth -p 3.11.9
# 進(jìn)入項(xiàng)目目錄
cd demo-unsloth
# 創(chuàng)建虛擬環(huán)境
uv venv
# 激活虛擬環(huán)境（Windows 下）
.venv\Scripts\activate
# 安裝 triton-windows，要求版本3.3.1.post19
uv pip install triton-windows==3.3.1.post19
# 安裝支持 CUDA 12.6 的 PyTorch 及其相關(guān)庫
uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

uv pip install unsloth

huggingface-cli download --resume-download unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit --local-dir ./ckpts/qwen2.5-1.5b-instruct-bnb-4bit

[
  {
    "Question": "傷寒一日，巨陽受之會出現(xiàn)什么癥狀？",
    "Response": "傷寒一日，巨陽受之會出現(xiàn)頭項(xiàng)痛、腰脊強(qiáng)的癥狀。因?yàn)榫揸栒?，諸陽之屬也，其脈連于風(fēng)府，為諸陽主氣，所以先受邪。 "
  },
  {
    "Question": "三陰受病，厥陰受之會出現(xiàn)什么癥狀",
    "Response": "三陰受病，若厥陰受之，厥陰脈循陰器而絡(luò)于肝，會出現(xiàn)煩滿而囊縮的癥狀。 若傷寒循常無變，十二日厥陰病衰，會出現(xiàn)囊縱、少腹微下的情況。 "
  },
]

# 導(dǎo)入必要的庫
from unsloth import FastLanguageModel, is_bfloat16_supported
from transformers import TrainingArguments, EarlyStoppingCallback
from trl import SFTTrainer
from datasets import load_dataset

# 模型配置參數(shù)
max_seq_length = 2048  # 模型處理文本的最大序列長度，支持長文本輸入

# 加載預(yù)訓(xùn)練模型和分詞器
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="ckpts/qwen2.5-1.5b-instruct-bnb-4bit",
    max_seq_length=max_seq_length,
    dtype=None,  # 自動檢測最佳數(shù)據(jù)類型（bfloat16或float16）
    load_in_4bit=True,  # 使用4bit量化加載模型，大幅減少顯存占用
)

“請先閱讀下方的任務(wù)指令，再結(jié)合提供的上下文信息，生成恰當(dāng)?shù)幕貜?fù)。### 指令：{} ### 上下文：{} ### 回復(fù)：{}”

# 定義訓(xùn)練數(shù)據(jù)格式化模板
# 使用中醫(yī)專家角色設(shè)定，專門針對《傷寒論》問答任務(wù)
train_prompt_style = """你是一位精通中醫(yī)理論的專家，特別擅長《傷寒論》的理論和實(shí)踐應(yīng)用。
請根據(jù)《傷寒論》的經(jīng)典理論，準(zhǔn)確回答以下問題。

### 問題:
{}

### 回答:
{}
"""

# 加載訓(xùn)練數(shù)據(jù)集
# 從JSON文件加載傷寒論問答數(shù)據(jù)集，包含問題和回答對
dataset = load_dataset("json", data_files="data/datasets-2025-08.json", split="train")  


def formatting_data(examples):
    """格式化數(shù)據(jù)集函數(shù)，將問答對轉(zhuǎn)換為訓(xùn)練格式
    
    將原始的問題和回答對格式化為模型訓(xùn)練所需的文本格式，
    添加角色設(shè)定和結(jié)構(gòu)化模板。

    Args:
        examples: 包含Question、Response字段的數(shù)據(jù)樣本字典

    Returns:
        dict: 包含格式化后文本的字典，鍵為"text"
    """
    questions = examples["Question"]
    responses = examples["Response"]
    texts = []
    for q, r in zip(questions, responses):
        # 使用模板格式化每個問答對，并添加結(jié)束標(biāo)記
        text = train_prompt_style.format(q, r) + tokenizer.eos_token
        texts.append(text)
    # print(f"數(shù)據(jù)集: {texts}")
    return {"text": texts}


# 應(yīng)用數(shù)據(jù)格式化函數(shù)
dataset = dataset.map(formatting_data, batched=True, num_proc=1)

# 數(shù)據(jù)集分割：80%用于訓(xùn)練，20%用于驗(yàn)證
# 使用固定隨機(jī)種子確保結(jié)果可復(fù)現(xiàn)
train_test_split = dataset.train_test_split(test_size=0.2, seed=3407)
train_dataset = train_test_split['train']
eval_dataset = train_test_split['test']

print(f"數(shù)據(jù)集加載完成 - 訓(xùn)練集: {len(train_dataset)} 條, 驗(yàn)證集: {len(eval_dataset)} 條")

# 添加LoRA（Low-Rank Adaptation）權(quán)重配置
# LoRA是一種高效的微調(diào)方法，只訓(xùn)練少量參數(shù)即可適應(yīng)新任務(wù)
model = FastLanguageModel.get_peft_model(
    model,
    r=32,  # LoRA矩陣的秩，值越大表達(dá)能力越強(qiáng)，但參數(shù)量也越多
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ],
    lora_alpha=64,  # LoRA縮放參數(shù)
    lora_dropout=0.1,  # LoRA dropout率，防止過擬合，值越小正則化越弱
    bias="none",  # 偏置項(xiàng)處理方式，"none"表示不訓(xùn)練偏置，節(jié)省參數(shù)
    use_gradient_checkpointing="unsloth",  # 使用Unsloth優(yōu)化的梯度檢查點(diǎn)，支持超長序列
    random_state=3407,  # 隨機(jī)種子，確保結(jié)果可復(fù)現(xiàn)
    use_rslora=False,  # 是否使用Rank Stabilized LoRA，當(dāng)前使用標(biāo)準(zhǔn)LoRA
    loftq_cnotallow=None,  # LoftQ配置，用于更精確的量化
)

# 創(chuàng)建SFT（Supervised Fine-Tuning）訓(xùn)練器
# 使用監(jiān)督學(xué)習(xí)方式微調(diào)模型，適用于問答任務(wù)
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,  # 訓(xùn)練數(shù)據(jù)集
    eval_dataset=eval_dataset,  # 驗(yàn)證數(shù)據(jù)集，用于評估模型性能
    dataset_text_field="text",  # 數(shù)據(jù)集中文本字段的名稱
    max_seq_length=max_seq_length,  # 最大序列長度
    dataset_num_proc=1,  # 數(shù)據(jù)處理進(jìn)程數(shù)，設(shè)為1避免緩存沖突
    packing=False,  # 是否使用序列打包，短序列時可設(shè)為True提升5倍訓(xùn)練速度
    callbacks=[
        EarlyStoppingCallback(early_stopping_patience=3, early_stopping_threshold=0.0005),  # 早停機(jī)制，防止過擬合
    ],
    args=TrainingArguments(
        # 批次大小配置
        per_device_train_batch_size=2,  # 每個GPU的訓(xùn)練批次大小
        per_device_eval_batch_size=2,  # 每個GPU的驗(yàn)證批次大小
        
        # 梯度累積配置
        gradient_accumulation_steps=8,  # 梯度累積步數(shù)，有效批次大小 = batch_size * gradient_accumulation_steps
        
        # 學(xué)習(xí)率配置
        warmup_ratio=0.15,  # 學(xué)習(xí)率預(yù)熱比例，前15%的步數(shù)用于預(yù)熱
        learning_rate=2e-5,  # 學(xué)習(xí)率，控制參數(shù)更新步長
        
        # 訓(xùn)練輪數(shù)和步數(shù)配置
        # max_steps = 200, # 最大訓(xùn)練步數(shù)（可選）
        num_train_epochs=5,  # 訓(xùn)練輪數(shù)，給模型充分學(xué)習(xí)機(jī)會
        
        # 精度配置
        fp16=not is_bfloat16_supported(),  # 是否使用16位浮點(diǎn)精度訓(xùn)練
        bf16=is_bfloat16_supported(),  # 是否使用bfloat16精度訓(xùn)練（更穩(wěn)定）
        
        # 日志和監(jiān)控配置
        logging_steps=2,  # 每2步記錄一次日志
        eval_steps=10,  # 每10步進(jìn)行一次驗(yàn)證評估
        eval_strategy="steps",  # 按步數(shù)進(jìn)行驗(yàn)證
        
        # 模型保存配置
        save_steps=20,  # 每20步保存一次模型檢查點(diǎn)
        save_strategy="steps",  # 按步數(shù)保存模型
        save_total_limit=5,  # 最多保存5個檢查點(diǎn)，節(jié)省存儲空間
        
        # 最佳模型配置
        load_best_model_at_end=True,  # 訓(xùn)練結(jié)束時自動加載最佳模型
        metric_for_best_model="eval_loss",  # 使用驗(yàn)證損失作為最佳模型指標(biāo)
        greater_is_better=False,  # 損失越小越好
        
        # 正則化配置
        weight_decay=0.001,  # 權(quán)重衰減，防止過擬合
        max_grad_norm=1.0,  # 梯度裁剪閾值，防止梯度爆炸
        
        # 學(xué)習(xí)率調(diào)度配置
        lr_scheduler_type="cosine",  # 使用余弦退火學(xué)習(xí)率調(diào)度器
        
        # 優(yōu)化器配置
        optim="adamw_8bit",  # 使用8位AdamW優(yōu)化器，節(jié)省顯存
        
        # 數(shù)據(jù)加載配置
        dataloader_num_workers=0,  # 數(shù)據(jù)加載器工作進(jìn)程數(shù)，設(shè)為0避免多進(jìn)程沖突
        
        # 輸出配置
        output_dir="outputs",  # 模型輸出和檢查點(diǎn)保存目錄
        
        # 隨機(jī)種子
        seed=3407,  # 隨機(jī)種子，確保結(jié)果可復(fù)現(xiàn)

    ),
)

# 開始訓(xùn)練
train_stats = trainer.train()
print(f"訓(xùn)練完成，訓(xùn)練損失: \n {train_stats}")

# 保存微調(diào)后的模型權(quán)重
# 只保存LoRA權(quán)重，原始模型權(quán)重保持不變
model.save_pretrained("ckpts/qwen2.5-1.5b-instruct-bnb-4bit-lora")
# 保存分詞器（tokenizer），以便后續(xù)加載和推理時使用
tokenizer.save_pretrained("ckpts/qwen2.5-1.5b-instruct-bnb-4bit-lora")
#==================================================================================

# model.save_pretrained_gguf("ckpts/qwen2.5-1.5b-instruct-bnb-4bit-gguf", tokenizer, quantization_method="q4_k_m")

==((====))==  Unsloth 2025.7.8: Fast Qwen3 patching. Transformers: 4.53.3.
   \\   /|    NVIDIA GeForce RTX 2050. Num GPUs = 1. Max memory: 4.0 GB. Platform: Windows.
O^O/ \_/ \    Torch: 2.7.1+cu126. CUDA: 8.6. CUDA Toolkit: 12.6. Triton: 3.3.1
\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.31.post1. FA2 = False]
 "-____-"     Free license: http://github.com/unslothai/unsloth
# 批注：系統(tǒng)配置信息
# - 使用RTX 2050顯卡，4GB顯存
# - 支持bfloat16精度訓(xùn)練
# - 使用Xformers優(yōu)化注意力機(jī)制

==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1
   \\   /|    Num examples = 353 | Num Epochs = 5 | Total steps = 115
O^O/ \_/ \    Batch size per device = 2 | Gradient accumulation steps = 8
\        /    Data Parallel GPUs = 1 | Total batch size (2 x 8 x 1) = 16
 "-____-"     Trainable parameters = 20,185,088 of 616,235,008 (3.28% trained)
# 批注：訓(xùn)練配置詳情
# - 總樣本數(shù)：353個
# - 訓(xùn)練輪數(shù)：5輪
# - 總步數(shù)：115步
# - 有效批次大?。?6（2×8×1）
# - 可訓(xùn)練參數(shù)：3.28%（20M/616M）

{'loss': 2.1017, 'grad_norm': 3.180413007736206, 'learning_rate': 1.995283421166614e-05, 'epoch': 1.05}
{'loss': 2.2702, 'grad_norm': 2.8307971954345703, 'learning_rate': 1.9869167087338908e-05, 'epoch': 1.14}
{'loss': 2.0511, 'grad_norm': 3.1757583618164062, 'learning_rate': 1.9744105246469264e-05, 'epoch': 1.23}
{'loss': 2.1853, 'grad_norm': 2.5234906673431396, 'learning_rate': 1.957817324187987e-05, 'epoch': 1.32}
{'loss': 1.8145, 'grad_norm': 2.9424679279327393, 'learning_rate': 1.937206705006344e-05, 'epoch': 1.32}

from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template
from transformers import TextStreamer
import torch

max_seq_length = 512


model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="ckpts/qwen2.5-1.5b-instruct-bnb-4bit-lora",
    max_seq_length=max_seq_length,
    dtype=None,
    load_in_4bit=True,
)

# 啟用unsloth的2x更快推理優(yōu)化
print("啟用推理優(yōu)化...")
FastLanguageModel.for_inference(model)
print("模型加載完成")


# 定義訓(xùn)練數(shù)據(jù)格式化字符串模板
train_prompt_style = """你是一位精通中醫(yī)理論的專家，特別擅長《傷寒論》的理論和實(shí)踐應(yīng)用。
請根據(jù)《傷寒論》的經(jīng)典理論，準(zhǔn)確回答以下問題。

### 問題:
{}
"""
question = '什么是太陽病？'

formatted_prompt = train_prompt_style.format(question)
print(f"格式化后的提示文本:\n-------------------\n{formatted_prompt}\n-------------------")
print(type(tokenizer))   # <class 'transformers.models.qwen2.tokenization_qwen2_fast.Qwen2TokenizerFast'>
inputs = tokenizer([formatted_prompt], return_tensors='pt', max_length=max_seq_length).to("cuda")

print(f"inputs: \n-------------------\n{inputs}\n-------------------")



# 生成回答（流式輸出）
with torch.no_grad():
    outputs = model.generate(
        # 輸入序列的token ID
        inputs['input_ids'], 
        # 注意力掩碼，指示哪些位置是有效輸入
        attention_mask=inputs['attention_mask'], 
        # 生成文本的最大長度限制
        max_length=max_seq_length,
        # 啟用KV緩存，加速生成過程
        use_cache=True,
        # 溫度參數(shù)，控制生成的隨機(jī)性
        # 值越低生成越確定，值越高生成越隨機(jī)
        temperature=0.8, 
        # 核采樣參數(shù)，只從累積概率達(dá)到90%的詞中選擇
        # 避免選擇概率極低的詞，提高生成質(zhì)量
        top_p=0.9,
        # 啟用采樣模式，而不是貪婪解碼
        # 貪婪解碼總是選擇概率最高的詞，容易產(chǎn)生重復(fù)
        do_sample=True,
        # 設(shè)置填充標(biāo)記ID，用于處理變長序列
        pad_token_id=tokenizer.eos_token_id,
        # 設(shè)置結(jié)束標(biāo)記ID，告訴模型何時停止生成
        eos_token_id=tokenizer.eos_token_id,
        # 重復(fù)懲罰參數(shù)，防止模型重復(fù)生成相同內(nèi)容
        # 1.1表示重復(fù)詞的生成概率降低10%
        repetition_penalty=1.1
    )

print(type(outputs))
print(outputs)
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("\n" + "=" * 80)
print(answer)
print("\n" + "=" * 80)
print("測試完成！")