自然语言处理中的文本分类领域详解：从传统方法到BERT实战代码

!pip install transformers
!pip install torch
!pip install pandas
!pip install sklearn

import pandas as pd
from sklearn.model_selection import train_test_split
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import Trainer, TrainingArguments
import torch

# 假设已有 IMDb 数据集
data = pd.read_csv("IMDB Dataset.csv")  # 数据集可在 https://ai.stanford.edu/~amaas/data/sentiment/ 获取
data = data.sample(5000)  # 选取5000条数据用于演示

# 文本和标签处理
data['label'] = data['sentiment'].apply(lambda x: 1 if x == 'positive' else 0)
train_texts, val_texts, train_labels, val_labels = train_test_split(data['review'].tolist(), data['label'].tolist(), test_size=0.2)

# 加载BERT分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=128)
val_encodings = tokenizer(val_texts, truncation=True, padding=True, max_length=128)

class IMDbDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

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

train_dataset = IMDbDataset(train_encodings, train_labels)
val_dataset = IMDbDataset(val_encodings, val_labels)

# 加载BERT模型
model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)

# 训练参数设置
training_args = TrainingArguments(
    output_dir='./results',          # 模型输出路径
    num_train_epochs=2,              # 训练轮数
    per_device_train_batch_size=8,   # 训练批次大小
    per_device_eval_batch_size=8,    # 验证批次大小
    warmup_steps=500,                # 学习率预热步数
    evaluation_strategy="epoch",     # 每个epoch验证
    logging_dir='./logs',            # 日志路径
)

# 使用Trainer API训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
)

# 模型训练
trainer.train()

# 模型评估
eval_results = trainer.evaluate()
print(f"Evaluation Results: {eval_results}")

# 保存模型
model.save_pretrained("./sentiment_model")
tokenizer.save_pretrained("./sentiment_model")

# 示例预测
inputs = tokenizer("This movie was fantastic!", return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
pred = torch.argmax(logits, dim=1).item()
print("Prediction:", "Positive" if pred == 1 else "Negative")