微调 BERT 模型：使用 Keras/TensorFlow 进行文本分类的完整指南

pip install tensorflow tensorflow-text transformers

import tensorflow as tf
import tensorflow_text as tf_text
from transformers import TFBertModel, BertTokenizer
import numpy as np

# 示例数据：包含正面和负面评论的列表
reviews = [
    "This movie was absolutely fantastic! I loved every minute of it.",
    "The plot was confusing and the acting was terrible.",
    "A masterpiece of modern cinema, highly recommended.",
    "I wasted two hours of my life watching this garbage.",
    "It was okay, nothing special, just average."
]
# 标签：1代表正面，0代表负面
labels = np.array([1, 0, 1, 0, 0])

model_name = 'bert-base-uncased'

# 加载分词器
tokenizer = BertTokenizer.from_pretrained(model_name)

def tokenize_and_pad(texts, tokenizer, max_len=128):
    # 使用tokenizer.batch_encode_plus进行批处理编码
    encoded_inputs = tokenizer.batch_encode_plus(
        texts,
        max_length=max_len,
        padding='max_length', # 使用max_length进行填充
        truncation=True,       # 截断超长的序列
        return_tensors='tf'    # 返回TensorFlow张量
    )
    
    # 提取所需的三种输入
    input_ids = encoded_inputs['input_ids']
    attention_mask = encoded_inputs['attention_mask']
    # token_type_ids 对于单句分类通常是全0，但这里我们让分词器处理
    token_type_ids = encoded_inputs.get('token_type_ids', tf.zeros_like(input_ids))
    
    return {
        'input_ids': input_ids,
        'attention_mask': attention_mask,
        'token_type_ids': token_type_ids
    }

# 准备数据
train_data = tokenize_and_pad(reviews, tokenizer)

# 1. 加载预训练的BERT模型
bert_model = TFBertModel.from_pretrained(model_name, include_pooling=True)

# 2. 构建自定义分类器模型
class BertClassifier(tf.keras.Model):
    def __init__(self, bert_model, num_classes):
        super(BertClassifier, self).__init__()
        # 确保bert_model是共享的基础层
        self.bert = bert_model
        # 添加一个Dropout层以防止过拟合
        self.dropout = tf.keras.layers.Dropout(0.1)
        # 添加一个密集层作为分类头
        self.classifier = tf.keras.layers.Dense(num_classes, activation='softmax')

    def call(self, inputs):
        # inputs 是一个字典，包含 input_ids, attention_mask, token_type_ids
        
        # BERT模型的输出是一个包含三个元素的元组：
        # (sequence_output, pooler_output, hidden_states, attentions)
        # pooler_output 是 [CLS] token 的输出（通常用于分类）
        bert_outputs = self.bert(inputs)
        
        # 我们只使用 pooler_output
        cls_output = bert_outputs.pooler_output
        
        # 传递给分类头
        x = self.dropout(cls_output)
        return self.classifier(x)

num_classes = 2 # 二元分类

bm_classifier = BertClassifier(bert_model, num_classes=num_classes)

# 冻结BERT层
bert_model.trainable = False

optimizer = tf.keras.optimizers.Adam(learning_rate=2e-5) # BERT微调通常使用较小的学习率
loss = tf.keras.losses.SparseCategoricalCrossentropy()
metrics = ['accuracy']

bm_classifier.compile(optimizer=optimizer, loss=loss, metrics=metrics)

# 检查模型结构 (可选)
# bm_classifier.build(input_shape=({k: (None,) for k in train_data.keys()}))
# bm_classifier.summary()

print("\n开始微调模型...")

# 训练模型
history = bm_classifier.fit(
    train_data,
    labels,
    epochs=3, # 对于演示，我们只训练少量轮次
    batch_size=2
)

print("\n模型训练完成。")

# 评估（在训练集上，仅用于演示）
loss, accuracy = bm_classifier.evaluate(train_data, labels)
print(f"\n训练集准确率: {accuracy*100:.2f}%")

# 进行预测
new_texts = [
    "This was truly awful and a waste of time.", 
    "I enjoyed the subtle complexity of the narrative."
]

pred_data = tokenize_and_pad(new_texts, tokenizer)
predictions = bm_classifier.predict(pred_data)

# 预测结果是概率分布，取最高概率的类别索引
predicted_classes = np.argmax(predictions, axis=1)

print("\n新文本预测结果 (0=负面, 1=正面):")
for text, pred in zip(new_texts, predicted_classes):
    print(f"  '{text[:30]}...' -> 类别: {pred}")