# Google Colab环境特殊安装流程
import os
if "COLAB_" in "".join(os.environ.keys()):
    # 安装基础依赖（不处理依赖关系）
    !pip install --no-deps bitsandbytes accelerate xformers==0.0.29.post3 peft trl==0.15.2 triton cut_cross_entropy unsloth_zoo
    # 安装NLP常用工具
    !pip install sentencepiece protobuf datasets huggingface_hub hf_transfer
    # 安装unsloth核心库
    !pip install --no-deps unsloth

    !pip install trl
else:
    # 本地环境直接安装
    !pip install unsloth

from unsloth import FastLanguageModel
import torch
from datasets import load_dataset, Dataset
import pandas as pd
from trl import SFTTrainer, SFTConfig
from unsloth.chat_templates import standardize_sharegpt
from transformers import TextStreamer

# 支持的4bit量化模型列表（部分）
fourbit_models = [
    "unsloth/Qwen3-1.7B-unsloth-bnb-4bit",  # Qwen 14B速度提升2倍
    "unsloth/Qwen3-14B-unsloth-bnb-4bit",
    "unsloth/Llama-3.1-8B",
    # 更多模型见Hugging Face Hub
]

# 加载Qwen3-14B模型（4bit量化）
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Qwen3-14B",
    max_seq_length=2048,    # 上下文长度，更长会消耗更多内存
    load_in_4bit=True,      # 启用4bit量化（低内存占用）
    load_in_8bit=False,     # 8bit量化（更高精度，内存占用翻倍）
    full_finetuning=False,  # 禁用全参数微调
    token=""                # 私有模型需提供访问令牌（https://huggingface.co/settings/tokens）
)


# fourbit_models = [
#     "unsloth/Qwen3-1.7B-unsloth-bnb-4bit", # Qwen 14B 2x faster
#     "unsloth/Qwen3-4B-unsloth-bnb-4bit",
#     "unsloth/Qwen3-8B-unsloth-bnb-4bit",
#     "unsloth/Qwen3-14B-unsloth-bnb-4bit",
#     "unsloth/Qwen3-32B-unsloth-bnb-4bit",

#     # 4bit dynamic quants for superior accuracy and low memory use
#     "unsloth/gemma-3-12b-it-unsloth-bnb-4bit",
#     "unsloth/Phi-4",
#     "unsloth/Llama-3.1-8B",
#     "unsloth/Llama-3.2-3B",
#     "unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit" # [NEW] We support TTS models!
# ] # More models at https://huggingface.co/unsloth

# model, tokenizer = FastLanguageModel.from_pretrained(
#     model_name = "unsloth/Qwen3-14B",
#     max_seq_length = 2048,   # Context length - can be longer, but uses more memory
#     load_in_4bit = True,     # 4bit uses much less memory
#     load_in_8bit = False,    # A bit more accurate, uses 2x memory
#     full_finetuning = False, # We have full finetuning now!
#     token = "",      # use one if using gated models
# )

# 添加LoRA适配器（仅更新1-10%参数）
model = FastLanguageModel.get_peft_model(
    model,
    r=32,                    # LoRA秩（建议值：8,16,32,64,128）
    target_modules=[         # 针对Qwen3模型的可训练模块
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    lora_alpha=32,           # LoRA alpha值（建议设为rank或rank*2）
    lora_dropout=0,          # LoRA dropout（0值为优化配置）
    bias="none",             # 偏置设置（"none"为优化选项）
    use_gradient_checkpointing="unsloth",  # 梯度检查点（减少30%显存）
    random_state=3407,       # 随机种子
    use_rslora=False,        # 禁用rank stabilized LoRA
    loftq_config=None        # LoftQ配置（None为默认）
)

!pip install -U datasets
# 加载数学推理数据集（AIMO竞赛级数据）
reasoning_dataset = load_dataset(
    "unsloth/OpenMathReasoning-mini", 
    split="cot",
    token=""  # 如需私有数据集请提供token
)

# 加载对话数据集（ShareGPT风格）
non_reasoning_dataset = load_dataset(
    "mlabonne/FineTome-100k", 
    split="train",
    token=""
)


# 查看推理数据集结构
reasoning_dataset

# 查看非推理数据集结构
non_reasoning_dataset

# 推理数据集转对话格式
def generate_conversation(examples):
    problems = examples["problem"]
    solutions = examples["generated_solution"]
    conversations = []
    for problem, solution in zip(problems, solutions):
        conversations.append([
            {"role": "user", "content": problem},
            {"role": "assistant", "content": solution}
        ])
    return {"conversations": conversations}

# 应用对话模板（不进行分词）
reasoning_conversations = tokenizer.apply_chat_template(
    reasoning_dataset.map(generate_conversation, batched=True)["conversations"],
    tokenize=False
)

# 对话数据集标准化处理
from unsloth.chat_templates import standardize_sharegpt
dataset = standardize_sharegpt(non_reasoning_dataset)
non_reasoning_conversations = tokenizer.apply_chat_template(
    dataset["conversations"],
    tokenize=False
)

# 查看转换后的第一个样本
reasoning_conversations[0]
# 查看转换后的第一个非推理样本
non_reasoning_conversations[0]
# 查看两个数据集的大小
print(len(reasoning_conversations))
print(len(non_reasoning_conversations))

# 设置数据混合比例（75%对话+25%推理）
chat_percentage = 0.75

# 按比例采样对话数据
non_reasoning_subset = pd.Series(non_reasoning_conversations).sample(
    n=int(len(reasoning_conversations) * (1.0 - chat_percentage)),
    random_state=2407
)

# 合并数据集
data = pd.concat([
    pd.Series(reasoning_conversations),
    non_reasoning_subset
])
data.name = "text"

# 转换为HuggingFace Dataset格式并打乱
combined_dataset = Dataset.from_pandas(pd.DataFrame(data))
combined_dataset = combined_dataset.shuffle(seed=3407)


# 查看数据集的基本信息
print(combined_dataset)


import pandas as pd

# 转换为pandas DataFrame以便更好地显示
df = pd.DataFrame(combined_dataset[:10])
display(df)

from trl import SFTTrainer, SFTConfig

# 初始化SFTTrainer
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=combined_dataset,
    eval_dataset=None,  # 可添加评估数据集
    args=SFTConfig(
        dataset_text_field="text",         # 数据集中的文本字段名
        per_device_train_batch_size=2,      # 每个设备的训练批次大小
        gradient_accumulation_steps=4,      # 梯度累积步数（模拟大批次）
        warmup_steps=5,                     # 预热步数
        max_steps=30,                       # 最大训练步数
        learning_rate=2e-4,                 # 学习率（长期训练可降至2e-5）
        logging_steps=1,                    # 日志记录间隔
        optim="adamw_8bit",                 # 优化器
        weight_decay=0.01,                  # 权重衰减
        lr_scheduler_type="linear",         # 学习率调度类型
        seed=3407,                          # 随机种子
        report_to="none"                    # 禁用实验追踪
    ),
)

# 显示训练前内存状态
gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024**3, 3)
max_memory = round(gpu_stats.total_memory / 1024**3, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")

# 开始训练
trainer_stats = trainer.train()

# 显示训练后内存与时间统计
used_memory = round(torch.cuda.max_memory_reserved() / 1024**3, 3)
used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory / max_memory * 100, 3)
lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.")
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")

# 构造对话消息
messages = [
    {"role": "user", "content": "Solve (x + 2)^2 = 0."}
]

# 应用对话模板（禁用思考模式）
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,  # 必须添加生成提示
    enable_thinking=False        # 禁用思考模式
)

# 文本生成（普通对话参数）
from transformers import TextStreamer
_ = model.generate(
    **tokenizer(text, return_tensors="pt").to("cuda"),
    max_new_tokens=256,           # 最大生成长度
    temperature=0.7, top_p=0.8, top_k=20,  # 普通对话参数
    streamer=TextStreamer(tokenizer, skip_prompt=True)
)

# 构造对话消息（同上）
messages = [
    {"role": "user", "content": "Solve (x + 2)^2 = 0."}
]

# 应用对话模板（启用思考模式）
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=True         # 启用思考模式
)

# 文本生成（推理模式参数）
_ = model.generate(
    **tokenizer(text, return_tensors="pt").to("cuda"),
    max_new_tokens=1024,          # 增加生成长度以获取完整推理过程
    temperature=0.6, top_p=0.95, top_k=20,  # 推理模式参数
    streamer=TextStreamer(tokenizer, skip_prompt=True)
)

# 本地保存LoRA适配器
model.save_pretrained("lora_model")
tokenizer.save_pretrained("lora_model")

# 上传到Hugging Face Hub
# model.push_to_hub("your_username/Qwen3-lora_model", token="your_token")
# tokenizer.push_to_hub("your_username/Qwen3-lora_model", token="your_token")

# 保存为4bit量化模型（for VLLM）
if True:
    model.save_pretrained_merged("model", tokenizer, save_method="merged_4bit_forced")
    # 上传到Hugging Face Hub
    model.push_to_hub_merged(
        "your_username/Qwen3-vLLM", 
        tokenizer, 
        save_method="merged_4bit_forced", 
        token="your_token"
    )

# 保存为q4_k_m格式GGUF（推荐）
if True:
    model.save_pretrained_gguf("model", tokenizer, quantization_method="q4_k_m")
    # 上传到Hugging Face Hub
    model.push_to_hub_gguf(
        "your_username/Qwen3-GGUF", 
        tokenizer, 
        quantization_method="q4_k_m", 
        token="your_token"
    )

# 保存到Google Drive
from google.colab import drive
drive.mount('/content/gdrive')
model.save_pretrained_gguf(
    "/content/gdrive/MyDrive/MyModel/model",
    tokenizer,
    quantization_method="q4_k_m"
)