# requirements: pip install datasets transformers torch numpy
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
import numpy as np

# 加载RTE（识别文本蕴含）
dataset = load_dataset("glue", "rte")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 先快速训练一个epoch得到基准模型
def tokenize_function(examples):
    return tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, padding="max_length", max_length=128)

tokenized_datasets = dataset.map(tokenize_function, batched=True)
train_dataset = tokenized_datasets["train"]
eval_dataset = tokenized_datasets["validation"]

model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)

training_args = TrainingArguments(
    output_dir="./rte_initial",
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    num_train_epochs=1,          # 仅训一个epoch用于检测噪声
    learning_rate=2e-5,
    evaluation_strategy="no",
    save_strategy="no",
    logging_steps=100,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)
trainer.train()

# 对训练集做预测，获取置信度
train_predictions = trainer.predict(train_dataset)  # 返回含loss和metrics的对象
# 对于两个类别，取softmax后最大值作为置信度
import torch
logits = torch.tensor(train_predictions.predictions)
probs = torch.nn.functional.softmax(logits, dim=-1)
confidence, predicted_labels = torch.max(probs, dim=-1)
one_hot_labels = np.eye(2)[train_dataset["label"]]  # 注意label是int，需转换
# 计算每条样本的交叉熵损失（也可以直接用logits和label计算）
loss = torch.nn.functional.cross_entropy(logits, torch.tensor(train_dataset["label"]), reduction="none")

# 设定阈值：比如选出loss最大的5%作为可疑样本
threshold = np.percentile(loss.numpy(), 95)
suspicious_indices = np.where(loss.numpy() > threshold)[0]
print(f"找到 {len(suspicious_indices)} 条可疑噪声样本（占总训练集 {len(train_dataset)} 的 {100*len(suspicious_indices)/len(train_dataset):.1f}%）")

# 创建清洗后的数据集
clean_indices = np.setdiff1d(np.arange(len(train_dataset)), suspicious_indices)
clean_train_dataset = train_dataset.select(clean_indices)
print(f"清洗后训练集大小: {len(clean_train_dataset)}")

# 可选：保存清洗后索引供后续使用
np.save("clean_rte_indices.npy", clean_indices)

import torch.nn as nn
import copy

def prune_bert_heads(model, keep_ratio=0.5):
    """
    按attention head的L1范数剪枝，保留top keep_ratio的head
    注意：仅适用于BERT-base-uncased（12层，12头）
    """
    config = model.config
    num_heads = config.num_attention_heads  # 12
    head_dim = config.hidden_size // num_heads  # 64

    for layer_idx, module in enumerate(model.bert.encoder.layer):
        attn = module.attention.self
        # 取出QKV的权重，合并计算重要性
        q_weight = attn.query.weight.data.view(num_heads, head_dim, -1)  # [12,64,768]
        k_weight = attn.key.weight.data.view(num_heads, head_dim, -1)
        v_weight = attn.value.weight.data.view(num_heads, head_dim, -1)
        # 每个head的重要性：L1范数之和
        importance = []
        for h in range(num_heads):
            imp = q_weight[h].abs().sum() + k_weight[h].abs().sum() + v_weight[h].abs().sum()
            importance.append(imp.item())
        importance = torch.tensor(importance)
        # 保留重要性最高的50% head
        threshold = torch.quantile(importance, 1 - keep_ratio)
        mask = importance >= threshold
        # 实际剪枝：将需要剪掉的head对应权重置0
        for h in range(num_heads):
            if not mask[h]:
                attn.query.weight.data[h*head_dim:(h+1)*head_dim, :] = 0
                attn.key.weight.data[h*head_dim:(h+1)*head_dim, :] = 0
                attn.value.weight.data[h*head_dim:(h+1)*head_dim, :] = 0
                # 也置零bias
                if attn.query.bias is not None:
                    attn.query.bias.data[h*head_dim:(h+1)*head_dim] = 0
                if attn.key.bias is not None:
                    attn.key.bias.data[h*head_dim:(h+1)*head_dim] = 0
                if attn.value.bias is not None:
                    attn.value.bias.data[h*head_dim:(h+1)*head_dim] = 0
        # 注意：output的dense也需要对应剪掉？为保持维度一致，我们这里只剪head本身（权重置0），输出维度不变，但后续attention计算时被剪掉的head输出为0，等价于剪枝。
        # 输出层的投影也需要处理（weight在output.dense，它输入是[hidden]），这里不处理也能工作，但计算量不减。真正的结构化剪枝应该同时修改out_proj维度。为简化demo，我们暂时只演示权重清零。
    return model

# 在清洗后的数据上微调，然后剪枝并再微调
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
# 先在清洗数据上微调2个epoch...（省略trainer代码）
# 然后剪枝
pruned_model = prune_bert_heads(model, keep_ratio=0.5)
# 再微调2个epoch恢复性能

# 保存时只保留模型参数
output_dir = "./final_model"
model.save_pretrained(output_dir)  # 默认只保存model state_dict（config+weights）
# 不保存optimizer

# 量化：使用PyTorch的dynamic quantization
import torch.quantization
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)
torch.save(quantized_model.state_dict(), "./model_int8.pt")
# 体积对比
import os
original_size = os.path.getsize(os.path.join(output_dir, "pytorch_model.bin")) / 1e6
quant_size = os.path.getsize("./model_int8.pt") / 1e6
print(f"原始模型大小: {original_size:.2f} MB; int8量化后: {quant_size:.2f} MB; 压缩比: {original_size/quant_size:.1f}x")