# %%
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import torch
import torch.nn as nn
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    get_cosine_schedule_with_warmup,
)
from tqdm import tqdm

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# %%
# 读取同文件夹下的对话 csv（列：问题、回答）
df = pd.read_csv("样本收集.csv", encoding="utf-8")
df = df.dropna(subset=["问题", "回答"]).reset_index(drop=True)
# 清洗空白字符：去首尾空白、合并连续空白为单个空格、替换中文全角空格
df["问题"] = df["问题"].astype(str).str.replace("\u3000", " ").str.replace(r"\s+", " ", regex=True).str.strip()
df["回答"] = df["回答"].astype(str).str.replace("\u3000", " ").str.replace(r"\s+", " ", regex=True).str.strip()
print(f"共 {len(df)} 条对话样本")
print(df.head())

# %%
# 加载 Qwen2.5-1.5B-Instruct，bf16 全量微调（单卡 16GB 够用）
model_name = "Qwen/Qwen2.5-1.5B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)
model.gradient_checkpointing_enable()
model.config.use_cache = False
model.to(device)

# 冻结前两层 decoder，节省显存
for i in range(2):
    for p in model.model.layers[i].parameters():
        p.requires_grad = False
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"可训练参数: {trainable}/{total} ({100*trainable/total:.2f}%)")
print(model)

# %%
# 对话数据集：把 (问题, 回答) 拼成提示文本，做因果语言建模
PROMPT_TEMPLATE = (
    "你是一名材料科学助手，请根据问题给出准确、专业的回答。\n"
    "问题：{q}\n"
    "回答：{a}{eos}"
)


class ChatDataset(Dataset):
    def __init__(self, dataframe, tokenizer, max_length=512):
        self.df = dataframe.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        q = str(self.df.loc[idx, "问题"])
        a = str(self.df.loc[idx, "回答"])
        text = PROMPT_TEMPLATE.format(q=q, a=a, eos=self.tokenizer.eos_token)
        tokens = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )
        input_ids = tokens["input_ids"].squeeze(0)
        attention_mask = tokens["attention_mask"].squeeze(0)
        # pad 位置用 -100 屏蔽损失
        labels = input_ids.clone()
        labels[labels == self.tokenizer.pad_token_id] = -100
        return input_ids, attention_mask, labels


def collate_fn(batch):
    input_ids = torch.stack([b[0] for b in batch])
    attention_mask = torch.stack([b[1] for b in batch])
    labels = torch.stack([b[2] for b in batch])
    return input_ids, attention_mask, labels


dataset = ChatDataset(df, tokenizer, max_length=512)
train_dataloader = DataLoader(
    dataset, batch_size=1, shuffle=True, num_workers=2, collate_fn=collate_fn
)
# 抽样查看一条
sample = dataset[0]
print("sample input_ids shape:", sample[0].shape)
print("decoded:", tokenizer.decode(sample[0], skip_special_tokens=True))

# %%
# 训练配置
epochs = 30
learning_rate = 2e-5
warmup_ratio = 0.05

optimizer = torch.optim.AdamW(
    [p for p in model.parameters() if p.requires_grad],
    lr=learning_rate,
    weight_decay=0.01,
)
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=int(warmup_ratio * epochs * len(train_dataloader)),
    num_training_steps=epochs * len(train_dataloader),
)

best_loss = float("inf")
save_dir = "./output"
os.makedirs(save_dir, exist_ok=True)

# %%
# 训练循环
for epoch in range(epochs):
    model.train()
    training_loss = 0.0
    progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch+1}/{epochs}")
    lens = 0
    for input_ids, attention_mask, labels in progress_bar:
        input_ids = input_ids.to(device)
        attention_mask = attention_mask.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

        training_loss += loss.item()
        lens += 1
        progress_bar.set_postfix({
            "loss": f"{loss.item():.4f}",
            "avg_loss": f"{training_loss / lens:.4f}",
            "lr": f"{scheduler.get_last_lr()[0]:.2e}",
        })

    scheduler.step()
    avg_train_loss = training_loss / lens
    print(f"Epoch {epoch+1} train_loss: {avg_train_loss:.4f}")

    # 仅记录最优 loss，不每个 epoch 保存（1.5B 整存太费时费空间）
    if avg_train_loss < best_loss:
        best_loss = avg_train_loss
        print(f"新最优 loss: {best_loss:.4f}")

# 训练结束保存整个最终模型
final_dir = os.path.join(save_dir, "final_model")
model.save_pretrained(final_dir)
tokenizer.save_pretrained(final_dir)
print(f"训练完成，最佳 loss: {best_loss:.4f}")
print(f"最终模型已保存到: {final_dir}")