import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from torch.optim import AdamW
from tqdm import tqdm


# 1. 定义专属的 PyTorch 文本数据集
class AITextDataset(Dataset):
    def __init__(self, csv_file, tokenizer, max_len=128):
        self.data = pd.read_csv(csv_file)
        self.tokenizer = tokenizer
        self.max_len = max_len

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

    def __getitem__(self, index):
        text = str(self.data.iloc[index, 0])
        label = int(self.data.iloc[index, 1])

        # 将汉字切成 token 序列
        encoding = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt',
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }


def train_text():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"💻 当前计算设备: {device}")
    if device.type == 'cpu':
        print("⚠️ 警告：当前使用 CPU 炼丹。NLP模型参数量巨大，这可能需要一些时间，请耐心等待...")

    print("正在加载预训练的中文 BERT 分词器与模型权重...")

    # 【核心修复】：换成了官方真实存在、最经典的 bert-base-chinese
    model_name = "bert-base-chinese"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
    model = model.to(device)

    print("正在封装数据...")
    dataset = AITextDataset('./data/text_dataset.csv', tokenizer, max_len=128)
    # 批量大小设为 8，防止 CPU 内存吃紧
    dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

    optimizer = AdamW(model.parameters(), lr=2e-5)

    epochs = 1

    print("\n🚀 --- 开始文本模型微调 ---")
    model.train()

    for epoch in range(epochs):
        progress_bar = tqdm(dataloader, desc=f"第 {epoch + 1}/{epochs} 轮", leave=True, colour='blue')
        running_loss = 0.0

        for batch in progress_bar:
            optimizer.zero_grad()

            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss

            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            progress_bar.set_postfix({'loss': f"{loss.item():.4f}"})

        print(f"✅ 第 {epoch + 1} 轮完成 | 平均 Loss: {running_loss / len(dataloader):.4f}")

    # 保存咱们微调后的专属大模型权重
    save_dir = "./finetuned_text_model"
    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)
    print(f"\n🎉 炼丹成功！专属的文本鉴别模型已保存在: {save_dir} 文件夹中。")


if __name__ == "__main__":
    train_text()