BERT 预训练模型深度解析

# 对比静态词向量 vs 上下文感知表示
from transformers import BertTokenizer, BertModel
import torch

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 同一个词 "bank" 在不同语境中的向量不同
sentences = [
    "I deposited money at the bank",
    "We sat on the river bank",
    "The pilot had to bank the plane",
]

for sent in sentences:
    inputs = tokenizer(sent, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 找到 "bank" 的位置
    tokens = tokenizer.tokenize(sent)
    bank_idx = tokens.index("bank") + 1  # +1 for [CLS]
    bank_vector = outputs.last_hidden_state[0, bank_idx]
    print(f"'bank' in '{sent}':")
    print(f"  向量均值: {bank_vector.mean().item():.6f}, L2范数: {bank_vector.norm().item():.4f}")

# Word2Vec 静态表示 vs BERT 上下文表示对比
import numpy as np

# 模拟 Word2Vec：同一个向量表示所有语境
word2vec_bank = np.random.randn(300)

sentences = [
    "river bank is muddy",
    "bank account balance",
    "plane started to bank",
]

print("=== Word2Vec (静态) ===")
for i, sent in enumerate(sentences):
    sim = np.dot(word2vec_bank, word2vec_bank) / (
        np.linalg.norm(word2vec_bank) ** 2
    )
    print(f"  '{sent}': 向量完全相同 (sim={sim:.4f})")

print("\n=== BERT (上下文感知) ===")
# 模拟 BERT 输出：不同语境下向量不同
bert_vectors = [np.random.randn(768) * (1 + 0.1 * i) for i in range(3)]
for i, (sent, vec) in enumerate(zip(sentences, bert_vectors)):
    sim = np.dot(bert_vectors[0], vec) / (
        np.linalg.norm(bert_vectors[0]) * np.linalg.norm(vec)
    )
    print(f"  '{sent}': 与句子0的余弦相似度 = {sim:.4f}")

# BERT 预训练+微调范式示意图代码
from transformers import BertForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset

# 第 1 步：预训练模型已经包含了通用语言理解能力
# （由 Google 在 Wikipedia + BooksCorpus 上训练）
model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=2  # 二分类
)
print(f"预训练参数数量: {model.num_parameters():,}")

# 第 2 步：用下游任务数据微调（只需少量标注数据）
# dataset = load_dataset('glue', 'mrpc')
# training_args = TrainingArguments(
#     output_dir='./results',
#     num_train_epochs=3,
#     per_device_train_batch_size=16,
#     learning_rate=2e-5,
# )
print("预训练权重已加载，可直接在下游任务上微调")

# 从零实现 Self-Attention
import numpy as np

def softmax(x):
    e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e_x / e_x.sum(axis=-1, keepdims=True)

def self_attention(Q, K, V):
    """缩放点积注意力"""
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)
    attn_weights = softmax(scores)
    output = attn_weights @ V
    return output, attn_weights

# 示例：3 个词，每个词 4 维
np.random.seed(42)
W_q = np.random.randn(4, 4)
W_k = np.random.randn(4, 4)
W_v = np.random.randn(4, 4)

embeddings = np.array([
    [1.0, 0.5, 0.2, 0.1],   # "我"
    [0.8, 1.2, 0.3, 0.5],   # "爱"
    [0.3, 0.6, 1.0, 0.8],   # "学习"
])

Q = embeddings @ W_q
K = embeddings @ W_k
V = embeddings @ W_v

output, weights = self_attention(Q, K, V)
print("注意力权重矩阵:")
print(np.round(weights, 3))
print("\n自注意力输出:")
print(np.round(output, 3))

# 多头注意力机制
import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    """多头注意力简化实现"""
    
    def __init__(self, d_model=768, n_heads=12):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x, mask=None):
        batch_size = x.size(0)
        
        # 线性投影 + 拆分成多头
        Q = self.W_q(x).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        # 缩放点积注意力
        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = torch.softmax(scores, dim=-1)
        
        output = (attn @ V).transpose(1, 2).contiguous()
        output = output.view(batch_size, -1, self.d_model)
        return self.W_o(output), attn

# BERT-base 参数分析
mha = MultiHeadAttention(d_model=768, n_heads=12)
print(f"多头注意力参数量: {sum(p.numel() for p in mha.parameters()):,}")
print(f"  每个头维度: {768 // 12} = 64")
print(f"  4 个线性层 × 768 × 768 × 2 ≈ {4 * 768 * 768:,}")

# BERT 编码器完整层结构
import torch
import torch.nn as nn

class BertEncoderLayer(nn.Module):
    """单层 BERT 编码器"""
    
    def __init__(self, d_model=768, n_heads=12, d_ff=3072, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
    
    def forward(self, x):
        # 子层 1: 多头自注意力 + 残差 + LayerNorm
        attn_out, _ = self.self_attn(x, x, x)
        x = self.ln1(x + attn_out)
        
        # 子层 2: 前馈网络 + 残差 + LayerNorm
        ffn_out = self.ffn(x)
        x = self.ln2(x + ffn_out)
        return x

# BERT-base: 12 层编码器
print("BERT-base 编码器配置:")
print(f"  隐藏维度: 768")
print(f"  注意力头数: 12")
print(f"  FFN 中间维度: 3072 (4×)")
print(f"  编码器层数: 12")
print(f"  最大序列长度: 512")
print(f"  总参数量: ~110M")

layer = BertEncoderLayer()
x = torch.randn(2, 128, 768)  # batch=2, seq=128, dim=768
out = layer(x)
print(f"\n输入形状: {x.shape} → 输出形状: {out.shape}")

# MLM 掩码策略实现
import numpy as np

def create_mlm_mask(tokens, mask_prob=0.15):
    """
    BERT 的 MLM 掩码策略
    80% -> [MASK], 10% -> 随机词, 10% -> 保持不变
    """
    masked_tokens = tokens.copy()
    labels = [-100] * len(tokens)  # -100 表示忽略
    
    # 随机选择 15% 的 token
    indices = np.random.choice(
        len(tokens), 
        size=max(1, int(len(tokens) * mask_prob)), 
        replace=False
    )
    
    vocab = ["the", "a", "is", "to", "of", "and", "in", "that", "it", "for"]
    
    for idx in indices:
        roll = np.random.random()
        if roll < 0.8:
            masked_tokens[idx] = "[MASK]"
        elif roll < 0.9:
            masked_tokens[idx] = np.random.choice(vocab)
        # else: 保持不变
        
        labels[idx] = tokens[idx]  # 标签始终是原词
    
    return masked_tokens, labels

# 演示
original = ["I", "love", "to", "learn", "natural", "language", "processing"]
np.random.seed(42)

for i in range(3):
    masked, labels = create_mlm_mask(original)
    print(f"原始: {' '.join(original)}")
    print(f"掩码: {' '.join(masked)}")
    print(f"标签: {[original[j] if labels[j] != -100 else '---' for j in range(len(original))]}")
    print()

# 使用 transformers 进行 MLM 推理
from transformers import pipeline

fill_mask = pipeline("fill-mask", model="bert-base-uncased")

# 测试不同的填空场景
examples = [
    "Paris is the capital of [MASK].",
    "The [MASK] is the largest planet in our solar system.",
    "Python is a popular [MASK] language.",
]

for text in examples:
    results = fill_mask(text)
    print(f"输入: {text}")
    for r in results[:3]:
        print(f"  {r['score']:.4f} → {r['sequence']}")
    print()

# 验证 80/10/10 策略的效果
# 在训练时，10% 的情况不替换为 [MASK]，这迫使模型
# 学会判断一个词是否需要被修改
print("MLM 损失计算：")
print("  只有被 mask 的位置计算损失")
print("  损失 = CrossEntropy(predictions[mask_positions], labels[mask_positions])")
print("  忽略位置 labels = -100")

# 分析 MLM 的学习过程
import torch
import torch.nn as nn

class MLMLossDemo:
    """演示 MLM 损失如何计算"""
    
    def __init__(self, vocab_size=30522):
        self.vocab_size = vocab_size
        self.criterion = nn.CrossEntropyLoss()
    
    def compute_loss(self, predictions, labels):
        """
        predictions: [batch, seq_len, vocab_size]
        labels: [batch, seq_len] (-100 表示忽略)
        """
        # 展平
        predictions = predictions.view(-1, self.vocab_size)
        labels = labels.view(-1)
        
        # 只计算被 mask 位置的损失
        mask = labels != -100
        loss = self.criterion(predictions[mask], labels[mask])
        return loss

demo = MLMLossDemo(vocab_size=1000)

# 模拟: batch=1, seq_len=10, 2个位置被 mask
predictions = torch.randn(1, 10, 1000)
labels = torch.tensor([[-100, -100, 42, -100, -100, -100, -100, 87, -100, -100]])

loss = demo.compute_loss(predictions, labels)
print(f"MLM Loss: {loss.item():.4f}")
print(f"有效位置: 2 / 10 (15% 掩码率的近似)")

# 掩码比例对训练的影响
print("\n掩码比例实验:")
for rate in [0.05, 0.10, 0.15, 0.25, 0.50]:
    n_masked = max(1, int(10 * rate))
    print(f"  mask_rate={rate:.0%}: {n_masked}/10 个位置被掩码")

# NSP 训练数据构造
import random

def create_nsp_pairs(sentences_by_doc, n_pairs=5):
    """构造 NSP 训练样本"""
    docs = list(sentences_by_doc.values())
    pairs = []
    
    for _ in range(n_pairs):
        doc = random.choice(docs)
        
        if random.random() < 0.5 and len(doc) >= 2:
            # 正样本：同一文档中的连续句子
            i = random.randint(0, len(doc) - 2)
            sent_a = doc[i]
            sent_b = doc[i + 1]
            label = 1  # IsNext
        else:
            # 负样本：不同文档的随机句子
            doc_a = random.choice(docs)
            doc_b = random.choice([d for d in docs if d != doc_a])
            sent_a = random.choice(doc_a)
            sent_b = random.choice(doc_b)
            label = 0  # NotNext
        
        pairs.append((sent_a, sent_b, label))
    
    return pairs

# 示例
corpus = {
    "doc1": ["BERT was published in 2018.", "It revolutionized NLP.", "Many variants followed."],
    "doc2": ["The weather is nice today.", "Let's go for a walk.", "The park is beautiful."],
    "doc3": ["Neural networks use backpropagation.", "Gradient descent updates weights.", "Learning rate is crucial."],
}

random.seed(42)
pairs = create_nsp_pairs(corpus, n_pairs=5)
for a, b, label in pairs:
    status = "IsNext" if label else "NotNext"
    print(f"[{status}] A: '{a}'")
    print(f"        B: '{b}'")
    print()

# NSP 分类头实现
import torch
import torch.nn as nn
from transformers import BertModel

class BertNSP(nn.Module):
    """BERT + NSP 分类头"""
    
    def __init__(self, model_name='bert-base-uncased'):
        super().__init__()
        self.bert = BertModel.from_pretrained(model_name)
        # NSP 分类头: [CLS] 输出 → 2 类
        self.classifier = nn.Sequential(
            nn.Linear(self.bert.config.hidden_size, 2),
        )
    
    def forward(self, input_ids, token_type_ids=None, attention_mask=None):
        outputs = self.bert(
            input_ids,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask
        )
        # [CLS] token 的输出 (第一个 token)
        cls_output = outputs.pooler_output
        logits = self.classifier(cls_output)
        return logits

# 使用示例
model = BertNSP()
tokenizer = torch.hub.load('huggingface/pytorch-transformers', 'tokenizer', 'bert-base-uncased')

# 模拟输入：两个句子
text_a = "The cat sat on the mat."
text_b = "It was very comfortable."

encoded = tokenizer(
    text_a, text_b,
    return_tensors="pt",
    padding=True,
    truncation=True,
    max_length=128
)

with torch.no_grad():
    logits = model(
        encoded["input_ids"],
        token_type_ids=encoded.get("token_type_ids"),
        attention_mask=encoded["attention_mask"]
    )
    probs = torch.softmax(logits, dim=-1)
    pred = torch.argmax(logits, dim=-1).item()
    
print(f"预测: {'IsNext' if pred == 1 else 'NotNext'}")
print(f"概率: IsNext={probs[0][1]:.4f}, NotNext={probs[0][0]:.4f}")

# 对比 NSP vs SOP (Sentence Order Prediction)
"""
NSP (原始 BERT):
  正样本: 同一文档的连续句子 A, B
  负样本: 不同文档的随机句子 A, B'
  问题: 模型可通过主题区分,而非学习句序逻辑

SOP (ALBERT 提出):
  正样本: 同一文档的连续句子 A, B
  负样本: 同一文档的连续句子 B, A (顺序调换)
  优点: 必须理解句序逻辑,无法靠主题作弊
"""

def create_sop_pairs(doc_sentences, n_pairs=3):
    """SOP 训练数据构造"""
    pairs = []
    for _ in range(n_pairs):
        idx = random.randint(0, len(doc_sentences) - 2)
        a, b = doc_sentences[idx], doc_sentences[idx + 1]
        
        if random.random() < 0.5:
            pairs.append((a, b, 1))  # 正确顺序
        else:
            pairs.append((b, a, 0))  # 调换了顺序
    
    return pairs

doc = [
    "First, we collect the data.",
    "Then, we preprocess it.",
    "Finally, we train the model.",
]

random.seed(42)
pairs = create_sop_pairs(doc)
for a, b, label in pairs:
    order = "正确→" if label else "←调换"
    print(f"[{order}] '{a}' → '{b}'")

# 文本分类微调
from transformers import BertForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset, DatasetDict

# 加载数据集（以 IMDB 情感分析为例）
# dataset = load_dataset("imdb")

# 加载预训练 BERT + 分类头
model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=2  # 正面/负面
)

# 训练参数
training_args = TrainingArguments(
    output_dir='./bert-imdb',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    learning_rate=2e-5,
    weight_decay=0.01,
    warmup_ratio=0.06,
    fp16=True,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
)

# trainer = Trainer(
#     model=model,
#     args=training_args,
#     train_dataset=dataset['train'],
#     eval_dataset=dataset['test'],
# )
# trainer.train()

print("分类微调配置:")
print(f"  模型: bert-base-uncased")
print(f"  分类头: Linear(768, 2)")
print(f"  学习率: 2e-5 (比预训练小 100 倍)")
print(f"  总可训练参数: {model.num_parameters():,}")

# NER 命名实体识别微调
from transformers import BertForTokenClassification
import torch

# 标签集合（BIO 标注方案）
ner_labels = [
    "O",                    # 非实体
    "B-PER", "I-PER",       # 人名
    "B-ORG", "I-ORG",       # 组织
    "B-LOC", "I-LOC",       # 地点
    "B-MISC", "I-MISC",     # 其他
]

model = BertForTokenClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=len(ner_labels)
)

# 模拟推理
from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
text = "John works at Google in New York"

inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
    predictions = torch.argmax(outputs.logits, dim=-1)

tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
pred_labels = [ner_labels[p] for p in predictions[0]]

print(f"文本: {text}")
print(f"Token		预测标签")
print("-" * 30)
for token, label in zip(tokens[1:-1], pred_labels[1:-1]):  # 去掉 [CLS] 和 [SEP]
    print(f"{token:<16} {label}")

# 问答系统微调（SQuAD 风格）
from transformers import BertForQuestionAnswering, BertTokenizer
import torch

model = BertForQuestionAnswering.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# 推理
context = "BERT was introduced by Google AI in 2018."
question = "Who introduced BERT?"

inputs = tokenizer(question, context, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    start_logits = outputs.start_logits
    end_logits = outputs.end_logits
    
    start_idx = torch.argmax(start_logits)
    end_idx = torch.argmax(end_logits)
    
    answer_tokens = inputs["input_ids"][0][start_idx:end_idx+1]
    answer = tokenizer.decode(answer_tokens)

print(f"问题: {question}")
print(f"上下文: {context}")
print(f"答案: {answer.strip()}")

# 文本相似度微调
from transformers import BertForSequenceClassification

sim_model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=2  # 相似/不相似
)

sent1 = "The cat is on the mat"
sent2 = "A feline sits on a rug"

inputs = tokenizer(sent1, sent2, return_tensors="pt")
with torch.no_grad():
    outputs = sim_model(**inputs)
    probs = torch.softmax(outputs.logits, dim=-1)

print(f"\n句子1: {sent1}")
print(f"句子2: {sent2}")
print(f"相似度概率: {probs[0][1]:.4f}")

# 对比不同 BERT 变体
from transformers import AutoModel, AutoTokenizer

variants = {
    "bert-base": "bert-base-uncased",
    "roberta-base": "roberta-base",
    "albert-base": "albert-base-v2",
    "distilbert": "distilbert-base-uncased",
    "deberta-base": "microsoft/deberta-base",
}

print(f"{'模型':<20} {'参数量':>12} {'层数':>6} {'隐藏维':>8}")
print("-" * 50)

configs = [
    ("bert-base",     "~110M", 12, 768),
    ("roberta-base",  "~125M", 12, 768),
    ("albert-base",   "~12M",  12, 768),  # 参数共享
    ("distilbert",    "~66M",  6,  768),
    ("deberta-base",  "~86M",  12, 768),
    ("bert-large",    "~340M", 24, 1024),
    ("albert-xxlarge","~18M",  12, 4096), # 参数共享
]

for name, params, layers, hidden in configs:
    print(f"{name:<20} {params:>12} {layers:>6} {hidden:>8}")

# 加载并测试推理速度
import time

text = "The quick brown fox jumps over the lazy dog."

for name, model_name in variants.items():
    try:
        tok = AutoTokenizer.from_pretrained(model_name)
        model = AutoModel.from_pretrained(model_name)
        
        inputs = tok(text, return_tensors="pt")
        
        # 预热
        with torch.no_grad():
            model(**inputs)
        
        # 计时
        start = time.time()
        for _ in range(10):
            with torch.no_grad():
                model(**inputs)
        elapsed = (time.time() - start) / 10
        
        print(f"{name}: {elapsed*1000:.1f}ms/次")
    except Exception as e:
        print(f"{name}: 加载失败 ({e})")

# 知识蒸馏：将 BERT 压缩为 DistilBERT
import torch
import torch.nn as nn
import torch.nn.functional as F

class DistilBertStudent(nn.Module):
    """简化的 DistilBERT 学生模型"""
    
    def __init__(self, teacher_hidden=768, student_hidden=768, n_layers=6):
        super().__init__()
        self.student_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=student_hidden,
                nhead=12,
                dim_feedforward=3072,
                dropout=0.1
            ) for _ in range(n_layers)
        ])
        # 投影层：学生隐藏维 → 教师隐藏维
        self.proj = nn.Linear(student_hidden, teacher_hidden)
    
    def forward(self, x):
        for layer in self.student_layers:
            x = layer(x)
        return x

class KnowledgeDistillationLoss(nn.Module):
    """蒸馏损失：CE + KL(软标签)"""
    
    def __init__(self, temperature=2.0, alpha=0.5):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss()
        self.kl_loss = nn.KLDivLoss(reduction="batchmean")
    
    def forward(self, student_logits, teacher_logits, hard_labels):
        # 硬标签损失（标准交叉熵）
        ce = self.ce_loss(student_logits, hard_labels)
        
        # 软标签损失（KL 散度）
        T = self.temperature
        soft_student = F.log_softmax(student_logits / T, dim=-1)
        soft_teacher = F.softmax(teacher_logits / T, dim=-1)
        kl = self.kl_loss(soft_student, soft_teacher) * (T ** 2)
        
        return self.alpha * ce + (1 - self.alpha) * kl

print("知识蒸馏原理:")
print("  教师: BERT-base (12 层, 110M 参数)")
print("  学生: DistilBERT (6 层, 66M 参数)")
print("  损失 = α·CE(hard) + (1-α)·KL(soft)")
print("  温度 T 控制软标签的平滑度")
print("  结果: 97% 性能, 60% 速度提升, 40% 大小减少")

# DeBERTa 解耦注意力机制
import torch
import torch.nn as nn
import math

class DisentangledSelfAttention(nn.Module):
    """DeBERTa 的解耦注意力（简化版）
    
    传统自注意力: A_ij = (q_i · k_j) / sqrt(d)
    解耦注意力: A_ij = (q_c_i · k_c_j + q_c_i · k_p_j + q_p_i · k_c_j) / sqrt(d)
    
    q_c: 内容查询向量
    q_p: 位置查询向量  
    k_c: 内容键向量
    k_p: 位置键向量
    """
    
    def __init__(self, d_model=768, n_heads=12):
        super().__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        
        # 内容投影
        self.W_qc = nn.Linear(d_model, d_model)
        self.W_kc = nn.Linear(d_model, d_model)
        # 位置投影
        self.W_qp = nn.Linear(d_model, d_model)
        self.W_kp = nn.Linear(d_model, d_model)
        # 值投影
        self.W_v = nn.Linear(d_model, d_model)
        
        self.pos_bias = nn.Embedding(512, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
    
    def forward(self, x, pos_ids):
        # 内容向量
        qc = self.W_qc(x)
        kc = self.W_kc(x)
        # 位置向量
        pos_emb = self.pos_bias(pos_ids)
        qp = self.W_qp(pos_emb)
        kp = self.W_kp(pos_emb)
        
        # 三项注意力分数（简化）
        v = self.W_v(x)
        # 实际实现需要考虑相对位置偏移
        output = x  # 简化
        return self.out_proj(output)

print("DeBERTa 解耦注意力 vs 传统自注意力:")
print("  传统: Q·K 统一计算")
print("  解耦: 内容-内容 + 内容-位置 + 位置-内容 三项分开")
print("  优势: 更好地建模相对位置关系")
print("  结果: SuperGLUE 超越人类水平")

# 完整的 BERT 文本分类微调流程
from transformers import (
    BertTokenizer,
    BertForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
)
from datasets import load_dataset
import numpy as np
import evaluate

# === 第 1 步：加载数据集 ===
# dataset = load_dataset("glue", "sst2")
# label_list = dataset["train"].features["label"].names
label_list = ["negative", "positive"]

# === 第 2 步：加载分词器和模型 ===
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=2,
    id2label={0: "negative", 1: "positive"},
    label2id={"negative": 0, "positive": 1},
)

# === 第 3 步：预处理 ===
# def tokenize(examples):
#     return tokenizer(examples["sentence"], truncation=True, max_length=128)
# 
# tokenized_datasets = dataset.map(tokenize, batched=True)
# data_collator = DataCollatorWithPadding(tokenizer)

# === 第 4 步：评估指标 ===
# metric = evaluate.load("accuracy")
# def compute_metrics(eval_pred):
#     logits, labels = eval_pred
#     predictions = np.argmax(logits, axis=-1)
#     return metric.compute(predictions=predictions, references=labels)

print("BERT 文本分类微调流程:")
print("  ✅ 加载预训练 BERT + 分类头")
print("  ✅ 设置训练参数（小学习率 + warmup）")
print("  ✅ 端到端微调所有参数")
print("  ✅ 早停防止过拟合")
print("  ✅ 保存最佳模型")

# 训练参数详解
from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./bert-classifier",
    
    # 训练轮数：2-4 轮足够，多了过拟合
    num_train_epochs=3,
    
    # Batch size：越大训练越快，但需要更多显存
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    
    # 学习率：比预训练小 100 倍
    learning_rate=2e-5,
    
    # Warmup：前 6% 的步骤学习率从 0 线性上升到峰值
    warmup_ratio=0.06,
    
    # 权重衰减：L2 正则化
    weight_decay=0.01,
    
    # 混合精度训练：加速 + 减少显存
    fp16=True,
    
    # 评估和保存策略
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    
    # 日志
    logging_steps=50,
    report_to="none",
)

# 学习率调度器可视化
print("学习率调度（Warmup + Linear Decay）：")
total_steps = 1000
warmup_steps = int(total_steps * 0.06)

for step in [0, 30, 60, 100, 200, 500, 800, 1000]:
    if step <= warmup_steps:
        lr = 2e-5 * (step / warmup_steps)
    else:
        lr = 2e-5 * (1 - (step - warmup_steps) / (total_steps - warmup_steps))
    bar = "█" * int(lr / 2e-5 * 20)
    print(f"  Step {step:4d}: {bar} ({lr:.1e})")

# 推理：使用微调后的模型
from transformers import pipeline

# 加载微调后的模型
classifier = pipeline(
    "sentiment-analysis",
    model="./bert-classifier",
    tokenizer="./bert-classifier",
    device=0  # GPU
)

# 批量推理
texts = [
    "This movie was absolutely brilliant! The acting was superb.",
    "Terrible waste of time. The plot made no sense.",
    "An average film with some good moments but also flaws.",
    "I couldn't stop watching. Best movie of the year!",
    "Boring, predictable, and overpriced.",
]

results = classifier(texts)

for text, result in zip(texts, results):
    label = result["label"]
    score = result["score"]
    emoji = "😊" if "positive" in label.lower() else "😞"
    print(f"{emoji} {label.upper():10s} ({score:.1%})")
    print(f"   {text[:60]}...")
    print()

# 不通过 pipeline 的底层推理
from transformers import BertTokenizer, BertForSequenceClassification
import torch

model = BertForSequenceClassification.from_pretrained("./bert-classifier")
tokenizer = BertTokenizer.from_pretrained("./bert-classifier")

text = "I love this product so much!"
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)

with torch.no_grad():
    outputs = model(**inputs)
    probs = torch.softmax(outputs.logits, dim=-1)
    pred = torch.argmax(outputs.logits, dim=-1)

print(f"文本: {text}")
print(f"预测: {model.config.id2label[pred.item()]}")
print(f"置信度: {probs[0][pred].item():.4f}")

# 高级技巧：梯度累积 + 多 GPU + 模型压缩
"""
当显存不足时使用梯度累积模拟大 batch size
"""

# 方法 1：梯度累积
# 等效 batch = per_device_batch × gradient_accumulation_steps × n_gpus
# 例如：16 × 4 × 2 = 128

training_args = TrainingArguments(
    output_dir="./bert-accumulated",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,  # 每 4 步更新一次
    learning_rate=2e-5,
    num_train_epochs=3,
    fp16=True,
)

# 方法 2：模型量化（推理加速）
from transformers import BertForSequenceClassification
import torch

model = BertForSequenceClassification.from_pretrained("./bert-classifier")

# INT8 量化
from optimum.int8 import INT8Quantizer

quantizer = INT8Quantizer.from_pretrained(model)
# quantized_model = quantizer.quantize()

print("模型压缩选项:")
print("  1. 知识蒸馏: BERT → DistilBERT (66M → 40% 大小)")
print("  2. 模型剪枝: 移除不重要的注意力头和 FFN 神经元")
print("  3. INT8 量化: 32-bit → 8-bit (4× 压缩)")
print("  4. 梯度累积: 模拟大 batch size 而不需要大显存")
print("  5. ONNX 导出: 转换为 ONNX 格式加速推理")

# 方法 3：ONNX 导出
# from optimum.onnxruntime import ORTModelForSequenceClassification
# model = ORTModelForSequenceClassification.from_pretrained(
#     "./bert-classifier",
#     export=True
# )
# model.save_pretrained("./bert-onnx")