Transformer 学习之路 - 法律标题文本分类实战
Transformer 学习之路 - 法律标题文本分类实战
Transformer 模型自 2017 年问世以来，迅速成为自然语言处理（NLP）领域的基石。它通过自注意力机制（Self-Attention）实现了对长距离依赖的高效建模，广泛应用于文本分类、机器翻译、问答系统等任务。本文将以法律标题文本分类为例，详细讲解如何使用 Transformer 技术解决实际问题。
1. 项目背景与目标
在法律领域，文本分类是一项基础但重要的任务。例如，将法律标题分类为“合同纠纷”、“知识产权”、“劳动法”等类别，有助于快速归档和检索。我们的目标是利用 Transformer 模型，实现法律标题的自动分类。
2. 环境准备与数据加载
首先，我们需要安装必要的 Python 库，并加载数据集。以下是代码示例：
! pip install transformers datasets evaluate peft accelerate gradio optimum sentencepiece
! pip install jupyterlab scikit-learn pandas matplotlib tensorboard nltk rouge
import evaluate
from datasets import load_dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer, DataCollatorWithPadding
# 加载数据集
dataset = load_dataset("csv", data_files="./data.csv", split="train", encoding="utf-8")
dataset = dataset.filter(lambda x: x["title"] is not None)
3. 数据预处理
在训练模型之前，我们需要对数据进行预处理。主要包括分词、标签映射和数据集分割。
# 获取所有的分类名称并创建映射
category_names = dataset["category_name"]
unique_category_names = list(set(category_names))
category_name_to_label = {cid: idx for idx, cid in enumerate(unique_category_names)}
# 分割训练集和测试集
datasets = dataset.train_test_split(test_size=0.1)
# 使用预训练的分词器进行分词
tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-macbert-base")
def preprocess_function(examples):
tokenized_inputs = tokenizer(examples["title"], max_length=128, truncation=True, padding="max_length")
tokenized_inputs["labels"] = [category_name_to_label[name] for name in examples["category_name"]]
return tokenized_inputs
tokenized_dataset = datasets.map(preprocess_function, batched=True, remove_columns=['category_id','__index_level_0__'])
4. 模型创建与配置
接下来，我们加载预训练的 Transformer 模型，并配置分类任务的输出类别数。
model = AutoModelForSequenceClassification.from_pretrained("hfl/chinese-macbert-base", num_labels=len(unique_category_names), cache_dir="./cache_dir")
# 配置模型的标签映射
id2label = {idx: cid for idx, cid in enumerate(unique_category_names)}
model.config.id2label = id2label
5. 评估函数与训练参数
为了监控模型的性能，我们定义了评估函数，并设置了训练参数。
acc_metric = evaluate.load("accuracy")
f1_metric = evaluate.load("f1")
def compute_metrics(eval_pred):
predictions, labels = eval_pred
predictions = predictions.argmax(axis=1)
acc = acc_metric.compute(predictions=predictions, references=labels)
f1 = f1_metric.compute(predictions=predictions, references=labels, average="macro")
acc.update(f1)
return acc
train_args = TrainingArguments(
output_dir="./checkpoints",
per_device_train_batch_size=32,
per_device_eval_batch_size=32,
logging_steps=10,
eval_strategy="epoch",
save_strategy="epoch",
save_total_limit=3,
learning_rate=2e-5,
weight_decay=0.01,
metric_for_best_model="f1",
load_best_model_at_end=True
)
6. 训练与评估
使用 Trainer 类进行模型训练，并在测试集上进行评估。
trainer = Trainer(
model=model,
args=train_args,
train_dataset=tokenized_dataset["train"],
eval_dataset=tokenized_dataset["test"],
data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
compute_metrics=compute_metrics
)
trainer.train()
trainer.evaluate(tokenized_dataset["test"])
7. 模型推理
训练完成后，我们可以使用模型对新数据进行推理。
from transformers import pipeline
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)
sen = "公司营业执照吊销的后果有哪些"
pipe(sen)
# 保存模型
model.save_pretrained(model_save_path)
8. 总结
通过本文的实战案例，我们深入了解了如何使用 Transformer 技术进行法律标题文本分类。从数据预处理到模型训练与评估，每一步都至关重要。Transformer 模型凭借其强大的建模能力，能够高效地处理复杂的文本分类任务，为法律领域的自动化处理提供了有力支持。
希望本文能帮助你更好地理解 Transformer 技术，并将其应用到实际项目中！