import pandas as pd
from sklearn.model_selection import train_test_split
from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments, default_data_collator
from datasets import Dataset
from tqdm import tqdm
import torch
from torch.nn import BCEWithLogitsLoss
from sklearn.metrics import accuracy_score, f1_score
import numpy as np

# 1. 加载数据集
import pandas as pd
import numpy as np

# 1. 读取数据
data = pd.read_csv('E:/our6.csv', header=None, names=['text', 'intents'])  # 假设您的数据集在这个路径

# 2. 处理标签，将其转化为列表
# 使用正则表达式提取双引号中的标签
import re

def extract_labels(intent_str):
    # 正则表达式提取双引号中的标签
    return re.findall(r'"(.*?)"', intent_str)

data['labels'] = data['intents'].apply(extract_labels)

# 3. 创建标签字典
unique_labels = set([label for sublist in data['labels'] for label in sublist])
label_to_id = {label: idx for idx, label in enumerate(unique_labels)}

# 显示标签字典以确认
print("标签字典:", label_to_id)

# 4. 将标签转化为标签 ID
data['label_ids'] = data['labels'].apply(lambda x: [label_to_id[label] for label in x if label in label_to_id])

# 5. 将标签标准化为固定长度的 one-hot 向量
def standardize_labels(label_ids, num_labels):
    labels = np.zeros(num_labels, dtype=np.float32)  # 确保数据类型为 float32
    for label_id in label_ids:
        if 0 <= label_id < num_labels:  # 确保 label_id 在有效范围内
            labels[label_id] = 1
        else:
            raise ValueError(f"Invalid label_id: {label_id}")
    return labels

num_labels = len(unique_labels)
data['label_one_hot'] = data['label_ids'].apply(lambda x: standardize_labels(x, num_labels))

# 显示数据框前几行以确认
print(data.head())
# 4. 切分数据集
train_data, val_data = train_test_split(data, test_size=0.2)

# 5. 创建 Hugging Face 数据集
train_dataset = Dataset.from_pandas(train_data)
val_dataset = Dataset.from_pandas(val_data)

# 6. 初始化 tokenizer 和 model
# 替换为您本地下载模型的路径
local_model_path = 'E:/data/moxing/bert/'  # 本地模型路径
tokenizer = BertTokenizer.from_pretrained(local_model_path)
model = BertForSequenceClassification.from_pretrained(local_model_path, num_labels=num_labels)

# 7. Tokenization
def tokenize_function(examples):
    tokenized = tokenizer(examples['text'], padding='max_length', truncation=True)
    tokenized['label_one_hot'] = [standardize_labels(x, num_labels) for x in examples['label_ids']]
    return tokenized

train_dataset = train_dataset.map(tokenize_function, batched=True)
val_dataset = val_dataset.map(tokenize_function, batched=True)

# 8. 自定义 DataCollator 处理多标签任务
class MultiLabelDataCollator:
    def __call__(self, features):
        labels = [f["label_one_hot"] for f in features]  # 使用 one-hot 编码的标签
        labels = torch.tensor(labels, dtype=torch.float32)  # 转换为张量
        batch = default_data_collator(features)
        batch["labels"] = labels
        return batch

# 9. 设置训练参数
training_args = TrainingArguments(
    output_dir='./results',
    eval_strategy='epoch',  # 修改为 eval_strategy
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,  # 每10步记录一次日志
)

# 10. 自定义评估指标计算函数
def compute_metrics(pred):
    logits, labels = pred
    preds = logits.sigmoid().round()  # 将预测值二值化
    acc = accuracy_score(labels, preds)
    f1 = f1_score(labels, preds, average='micro')  # 根据需要选择'micro'或'macro'
    return {"accuracy": acc, "f1": f1}

# 11. 创建 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    data_collator=MultiLabelDataCollator(),  # 使用自定义的多标签 DataCollator
    compute_metrics=compute_metrics,  # 计算多标签分类的评估指标
)

# 12. 自定义训练过程，包含进度条显示
def train_with_progress_bar(trainer):
    total_steps = trainer.args.num_train_epochs * len(train_dataset) // trainer.args.per_device_train_batch_size
    progress_bar = tqdm(total=total_steps, desc="Training Progress")

    for epoch in range(trainer.args.num_train_epochs):
        for step, batch in enumerate(trainer.get_train_dataloader()):
            trainer.training_step(batch, step)
            progress_bar.update(1)

            # 每个step结束时执行验证
            if (step + 1) % trainer.args.logging_steps == 0:
                trainer.evaluate()

        # 每个epoch结束时执行验证
        trainer.evaluate()

    progress_bar.close()

# 13. 开始训练
train_with_progress_bar(trainer)

# 14. 保存模型
model.save_pretrained('./multi_intent_model')
tokenizer.save_pretrained('./multi_intent_model')仔细阅读代码，阅读完毕后回答ok，之后我会提问问题