from unsloth import FastLanguageModel
from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments
import torch

# 加载模型 这个模型是我自己的路径 根据自己模型的不同 填写不同的路径
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="/mnt/workspace/DeepSeek-R1-Distill-Qwen-7B",
    max_seq_length=2048,
    dtype=torch.bfloat16,
    load_in_4bit=False,
    device_map="cuda",
    offload_folder=None,
)
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = model.config.eos_token_id

# 加载数据集 数据集一样也是自己的路径 下载路径：#   https://huggingface.co/datasets/FreedomIntelligence/medical-o1-reasoning-SFT
dataset = load_dataset("json", data_files="/mnt/workspace/medical-o1-reasoning-SFT/medical_o1_sft.json", split="train")
dataset = dataset.select(range(500))
print(f"数据集大小: {len(dataset)}")

# 定义提示模板
train_prompt_style = """Below is an instruction that describes a task, paired with an input that provides further context. 
Write a response that appropriately completes the request. 
Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.

### Instruction:
You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning. 
Please answer the following medical question. 

### Question:
{}

### Response:
<think>
{}
</think>
{}"""

EOS_TOKEN = tokenizer.eos_token

# 定义格式化函数
def formatting_prompts_func(examples):
    inputs = examples["Question"]
    cots = examples["Complex_CoT"]
    outputs = examples["Response"]
    texts = []
    for input, cot, output in zip(inputs, cots, outputs):
        text = train_prompt_style.format(input, cot, output) + EOS_TOKEN
        texts.append(text)
    return {"text": texts}

# 格式化数据集
dataset = dataset.map(formatting_prompts_func, batched=True)
print("第一条格式化数据：", dataset["text"][0])

# 配置 LoRA
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    lora_dropout=0,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=3407,
)

# 训练参数
training_args = TrainingArguments(
    output_dir="/mnt/workspace/medical-o1-finetuned",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    num_train_epochs=1,
    learning_rate=2e-4,
    fp16=False,
    bf16=True,
    logging_steps=10,
    save_steps=100,
    max_grad_norm=0.3,
    warmup_steps=10,
)

# 初始化 trainer
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=2048,
    args=training_args,
)

# 开始微调
trainer.train()

# 保存模型
model.save_pretrained("/mnt/workspace/medical-o1-finetuned/model")
tokenizer.save_pretrained("/mnt/workspace/medical-o1-finetuned/tokenizer")