从零预训练领域BERT:法律文本理解的实战记录

2022-09-05 | 预训练 BERT NLP

通用BERT在法律文本上表现不佳,我们花了3个月从零预训练了一个法律领域BERT。这篇文章记录完整的数据准备、训练配置、踩坑经历。最终在合同条款分类任务上F1从78%提升到89%。

为什么需要领域预训练

2022年初,我们在做合同智能审核项目。用bert-base-chinese做条款分类,F1只有78%。分析了bad case后发现:

"甲方"、"乙方"、"违约金"这些词在通用语料中很少出现
法律文本的句式结构和日常文本差异很大
很多法律术语被拆成了[UNK]

我们决定在法律语料上继续预训练,让模型"学会"法律语言。

数据收集与清洗

数据来源

来源	数量	清洗后	说明
裁判文书网	800万篇	2.1GB	去除格式、当事人信息
法律法规	50万条	0.3GB	现行有效的法规
合同模板	10万份	0.4GB	脱敏处理
法律问答	200万对	0.2GB	律师咨询记录
总计	-	3.0GB	纯文本

清洗脚本

import re
from typing import Iterator

def clean_legal_text(text: str) -> str:
    """法律文本清洗"""
    # 1. 去除案号等敏感信息
    text = re.sub(r'\(\d{4}\)[^号]*号', '', text)
    
    # 2. 去除当事人姓名(脱敏)
    text = re.sub(r'(原告|被告|上诉人|被上诉人)[：:]\s*[\u4e00-\u9fa5]{2,4}', 
                  r'\1：XXX', text)
    
    # 3. 去除多余空白
    text = re.sub(r'\s+', ' ', text)
    
    # 4. 去除过短的段落(通常是格式残留)
    paragraphs = text.split('\n')
    paragraphs = [p.strip() for p in paragraphs if len(p.strip()) > 20]
    
    return '\n'.join(paragraphs)

def process_file(filepath: str) -> Iterator[str]:
    """处理单个文件"""
    with open(filepath, 'r', encoding='utf-8') as f:
        for line in f:
            cleaned = clean_legal_text(line)
            if len(cleaned) > 50:  # 过滤过短文本
                yield cleaned

⚠️ 踩坑1:数据质量比数量重要
最初我们收集了10GB数据,但包含大量OCR错误和格式问题。训练出的模型在下游任务上甚至不如bert-base。清洗到3GB后效果才上来。

词表扩展

为什么要扩展词表

bert-base-chinese的词表有21128个token,但很多法律术语不在里面:

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

# 看看法律术语是怎么被切分的
print(tokenizer.tokenize("违约金"))    # ['违', '约', '金'] - OK
print(tokenizer.tokenize("担保物权"))  # ['担', '保', '物', '权'] - OK
print(tokenizer.tokenize("善意取得"))  # ['善', '意', '取', '得'] - OK
print(tokenizer.tokenize("抵押权"))    # ['抵', '[UNK]', '权'] - 问题!

构建领域词表

from collections import Counter
import jieba

def build_domain_vocab(corpus_files: list, min_freq: int = 100) -> list:
    """从语料中提取高频领域词汇"""
    word_freq = Counter()
    
    for filepath in corpus_files:
        with open(filepath, 'r') as f:
            for line in f:
                words = jieba.cut(line.strip())
                word_freq.update(words)
    
    # 过滤:长度2-4,频率>=min_freq,且不在原词表中
    base_tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
    base_vocab = set(base_tokenizer.vocab.keys())
    
    domain_words = []
    for word, freq in word_freq.most_common():
        if (freq >= min_freq and 
            2 <= len(word) <= 4 and 
            word not in base_vocab):
            domain_words.append(word)
    
    return domain_words[:5000]  # 最多添加5000个词

# 扩展词表
new_words = build_domain_vocab(corpus_files)
# ['诉讼费', '管辖权', '被执行人', '申请执行', '强制执行', ...]

💡 词表扩展策略:
最终我们添加了3872个法律领域词汇,词表从21128扩展到24000。embedding层需要resize,新词的embedding随机初始化。

预训练配置

硬件环境

4 × NVIDIA V100 32GB
64核CPU,256GB内存
NVMe SSD存储

训练参数

# train_config.yaml
model:
  vocab_size: 24000
  hidden_size: 768
  num_hidden_layers: 12
  num_attention_heads: 12
  max_position_embeddings: 512

training:
  # 从bert-base-chinese初始化
  init_checkpoint: "bert-base-chinese"
  
  # 数据
  train_batch_size: 64  # per GPU
  max_seq_length: 512
  
  # 优化器
  learning_rate: 2e-5  # 继续预训练用较小lr
  weight_decay: 0.01
  warmup_ratio: 0.1
  
  # 训练量
  num_train_epochs: 3
  save_steps: 5000
  
  # MLM配置
  mlm_probability: 0.15
  
  # 混合精度
  fp16: true

⚠️ 踩坑2:学习率不能太大
最初用1e-4的学习率,训练几个epoch后loss反而上升。后来发现继续预训练和从头训练不同,需要用较小的学习率(2e-5)保护已学到的知识。

训练脚本

import torch
from transformers import (
    BertForMaskedLM, 
    BertTokenizer,
    DataCollatorForLanguageModeling,
    Trainer, 
    TrainingArguments
)
from datasets import load_dataset

def main():
    # 加载tokenizer(已扩展词表)
    tokenizer = BertTokenizer.from_pretrained('./legal-vocab')
    
    # 加载模型,resize embedding
    model = BertForMaskedLM.from_pretrained('bert-base-chinese')
    model.resize_token_embeddings(len(tokenizer))
    
    # 加载数据
    dataset = load_dataset('text', data_files={'train': 'legal_corpus.txt'})
    
    def tokenize_function(examples):
        return tokenizer(
            examples['text'],
            padding='max_length',
            truncation=True,
            max_length=512,
            return_special_tokens_mask=True
        )
    
    tokenized_dataset = dataset.map(
        tokenize_function, 
        batched=True,
        remove_columns=['text']
    )
    
    # MLM数据处理
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=True,
        mlm_probability=0.15
    )
    
    # 训练参数
    training_args = TrainingArguments(
        output_dir='./legal-bert',
        num_train_epochs=3,
        per_device_train_batch_size=64,
        learning_rate=2e-5,
        warmup_ratio=0.1,
        weight_decay=0.01,
        save_steps=5000,
        logging_steps=100,
        fp16=True,
        dataloader_num_workers=8,
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=tokenized_dataset['train'],
    )
    
    trainer.train()
    trainer.save_model('./legal-bert-final')

if __name__ == '__main__':
    main()

训练过程监控

Loss曲线

训练步数	MLM Loss	说明
0	2.85	初始(bert-base)
10000	2.12	快速下降
50000	1.65	下降变缓
100000	1.42	接近收敛
150000	1.38	基本收敛

完整训练耗时约5天(120小时),电费约800元。

下游任务评测

任务1:合同条款分类

模型	Precision	Recall	F1
bert-base-chinese	76.2%	80.1%	78.1%
RoBERTa-wwm	79.5%	82.3%	80.9%
Legal-BERT(ours)	87.3%	90.8%	89.0%

任务2:法律命名实体识别

模型	F1
bert-base-chinese	82.5%
Legal-BERT	88.7%

✅ 核心收益: 领域预训练在我们的任务上带来了8-11个点的F1提升,投入产出比很高。

经验总结

值得做的

数据清洗:花2周时间清洗数据是值得的,脏数据会毁掉整个模型
词表扩展:领域术语加入词表后,tokenization质量明显提升
继续预训练而非从头训练:从bert-base初始化,训练成本低很多

可以跳过的

NSP任务:我们去掉了NSP,只做MLM,效果没有明显差异
过长的训练:3个epoch后loss基本不降了,更多epoch收益很小

未解决的问题

新增词汇的embedding初始化:随机初始化 vs 近义词平均,哪个更好?
领域迁移:在法律语料上训练后,通用能力是否下降?

参考资料

Don't Stop Pretraining: Adapt Language Models to Domains and Tasks (ACL 2020)
How to Train BERT with Academic Resources
Chinese-BERT-wwm技术报告

更新记录:
2022-09-05: 初版发布
2023-02-10: 补充词表扩展细节
2023-08-20: 模型已在HuggingFace开源 (legal-bert-base-chinese)