NLP - 基于bert预训练模型的文本多分类示例

项目说明

项目名称

基于DistilBERT的标题多分类任务

项目概述

本项目旨在使用DistilBERT模型对给定的标题文本进行多分类任务。项目包括从数据处理、模型训练、模型评估到最终的API部署。该项目采用模块化设计，以便于理解和维护。

项目结构

.
├── bert_data
│   ├── train.txt
│   ├── dev.txt
│   └── test.txt
├── saved_model
├── results
├── logs
├── data_processing.py
├── dataset.py
├── training.py
├── app.py
└── main.py

文件说明

1. bert_data/：存放训练集、验证集和测试集的数据文件。

   • train.txt
   • dev.txt
   • test.txt

2. saved_model/：存放训练好的模型和tokenizer。

3. results/：存放训练结果。

4. logs/：存放训练日志。

5. data_processing.py：数据处理模块，负责读取和预处理数据。

6. dataset.py：数据集类模块，定义了用于训练和评估的数据集类。

7. training.py：模型训练模块，定义了训练和评估模型的过程。

8. app.py：模型部署模块，使用FastAPI创建API服务。

9. main.py：主脚本，运行整个流程，包括数据处理、模型训练和部署。

数据集数据规范

为了确保数据处理和模型训练的顺利进行，请按照以下规范准备数据集文件。每个文件包含的标题和标签分别使用制表符（\t）分隔。以下是一个示例数据集的格式。

数据文件格式

数据文件应为纯文本文件，扩展名为.txt，文件内容的每一行应包含一个文本标题和一个对应的分类标签，用制表符分隔。数据文件不应包含表头。

数据示例

探索神秘的海底世界    7
如何在家中制作美味披萨    2
全球气候变化的原因和影响    1
最新的智能手机评测    8
健康饮食：如何搭配均衡的膳食    5
最受欢迎的电影和电视剧推荐    3
了解宇宙的奥秘：天文学入门    0
如何种植和照顾多肉植物    9
时尚潮流：今年夏天的必备单品    6
如何有效管理个人财务    4

注意事项

• 标签规范：确保每个标题文本的标签是一个整数，表示类别。
• 文本编码：确保数据文件使用UTF-8编码，避免中文字符乱码。
• 数据一致性：确保训练、验证和测试数据格式一致，便于数据加载和处理。

通过以上规范和示例数据文件创建方法，可以确保数据文件符合项目需求，并顺利进行数据处理和模型训练。

模块说明

1. 数据处理模块 (data_processing.py)

功能：读取数据文件并进行预处理。

• load_data(file_path): 读取指定路径的数据文件，并返回一个包含文本和标签的数据框。
• tokenize_data(data, tokenizer, max_length=128): 使用BERT的tokenizer对数据进行tokenize处理。
• main(): 加载数据、tokenize数据并返回处理后的数据。

2. 数据集类模块 (dataset.py)

功能：定义数据集类，便于模型训练。

• TextDataset: 将tokenized数据和标签封装成PyTorch的数据集格式，便于Trainer进行训练和评估。

3. 模型训练模块 (training.py)

功能：定义训练和评估模型的过程。

• train_model(): 加载数据和tokenizer，创建数据集，加载模型，设置训练参数，定义Trainer，训练和评估模型，保存训练好的模型和tokenizer。

4. 模型部署模块 (app.py)

功能：使用FastAPI进行模型部署。

• predict(item: Item): 接收POST请求的文本输入，使用训练好的模型进行预测并返回分类结果。
• FastAPI应用启动配置。

5. 主脚本 (main.py)

功能：运行整个流程，包括数据处理、模型训练和部署。

• main(): 运行模型训练流程，并输出训练完成的提示。

运行步骤

1. 安装依赖

pip install pandas torch transformers fastapi uvicorn scikit-learn

2. 数据处理

确保bert_data文件夹下包含train.txt、dev.txt和test.txt文件，每个文件包含文本和标签，使用制表符分隔。

3. 训练模型

运行main.py脚本，进行数据处理和模型训练：

python main.py

训练完成后，模型和tokenizer将保存在saved_model文件夹中。

4. 部署模型

运行app.py脚本，启动API服务：

uvicorn app:app --reload

服务启动后，可以通过POST请求访问预测接口，进行文本分类预测。

示例请求

curl -X POST "http://localhost:8000/predict" -H "Content-Type: application/json" -d '{"text": "你的文本"}'

返回示例：

{"prediction": 3
}

注意事项

• 确保数据文件格式正确，每行包含一个文本和对应的标签，使用制表符分隔。
• 调整训练参数（如batch size和训练轮数）以适应不同的GPU配置。
• 使用nvidia-smi监控显存使用，避免显存溢出。

项目代码

1. 数据处理模块

功能：读取数据文件并进行预处理。

# data_processing.py
import pandas as pd
from transformers import DistilBertTokenizerdef load_data(file_path):data = pd.read_csv(file_path, delimiter='\t', header=None)data.columns = ['text', 'label']return datadef tokenize_data(data, tokenizer, max_length=128):encodings = tokenizer(list(data['text']), truncation=True, padding=True, max_length=max_length)return encodingsdef main():# 加载Tokenizertokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-chinese')# 加载数据train_data = load_data('./bert_data/train.txt')dev_data = load_data('./bert_data/dev.txt')test_data = load_data('./bert_data/test.txt')# Tokenize数据train_encodings = tokenize_data(train_data, tokenizer)dev_encodings = tokenize_data(dev_data, tokenizer)test_encodings = tokenize_data(test_data, tokenizer)return train_encodings, dev_encodings, test_encodings, train_data['label'], dev_data['label'], test_data['label']if __name__ == "__main__":main()

2. 数据集类模块

功能：定义数据集类，便于模型训练。

# dataset.py
import torchclass TextDataset(torch.utils.data.Dataset):def __init__(self, encodings, labels):self.encodings = encodingsself.labels = labelsdef __getitem__(self, idx):item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}item['labels'] = torch.tensor(self.labels[idx])return itemdef __len__(self):return len(self.labels)

3. 模型训练模块

功能：定义训练和评估模型的过程。

# training.py
import torch
from transformers import DistilBertForSequenceClassification, Trainer, TrainingArguments
from dataset import TextDataset
import data_processingdef train_model():# 加载数据和tokenizertrain_encodings, dev_encodings, test_encodings, train_labels, dev_labels, test_labels = data_processing.main()# 创建数据集train_dataset = TextDataset(train_encodings, train_labels)dev_dataset = TextDataset(dev_encodings, dev_labels)test_dataset = TextDataset(test_encodings, test_labels)# 加载DistilBERT模型model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-chinese', num_labels=10)model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))# 设置训练参数training_args = TrainingArguments(output_dir='./results',          # 输出结果目录num_train_epochs=3,              # 训练轮数per_device_train_batch_size=16,  # 训练时每个设备的批量大小per_device_eval_batch_size=64,   # 验证时每个设备的批量大小warmup_steps=500,                # 训练步数weight_decay=0.01,               # 权重衰减logging_dir='./logs',            # 日志目录fp16=True,                       # 启用混合精度训练)# 定义Trainertrainer = Trainer(model=model,                         # 预训练模型args=training_args,                  # 训练参数train_dataset=train_dataset,         # 训练数据集eval_dataset=dev_dataset             # 验证数据集)# 训练模型trainer.train()# 评估模型eval_results = trainer.evaluate()print(eval_results)# 保存模型trainer.save_model('./saved_model')tokenizer = trainer.tokenizertokenizer.save_pretrained('./saved_model')if __name__ == "__main__":train_model()

4. 模型部署模块

功能：使用FastAPI进行模型部署。

# app.py
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import torchapp = FastAPI()# 加载模型和tokenizer
model = DistilBertForSequenceClassification.from_pretrained('./saved_model')
tokenizer = DistilBertTokenizer.from_pretrained('./saved_model')class Item(BaseModel):text: str@app.post("/predict")
def predict(item: Item):inputs = tokenizer(item.text, return_tensors="pt", max_length=128, padding='max_length', truncation=True)outputs = model(**inputs)prediction = torch.argmax(outputs.logits, dim=1)return {"prediction": prediction.item()}if __name__ == "__main__":import uvicornuvicorn.run(app, host="0.0.0.0", port=8000)

5. 主脚本

功能：运行整个流程，包括数据处理、模型训练和部署。

# main.py
import trainingdef main():# 训练模型training.train_model()print("模型训练完成并保存。")if __name__ == "__main__":main()

其他：客户端调用案例

# client.py
import requestsdef predict(text):url = "http://localhost:8000/predict"payload = {"text": text}headers = {"Content-Type": "application/json"}response = requests.post(url, json=payload, headers=headers)if response.status_code == 200:prediction = response.json()return predictionelse:print(f"Error: {response.status_code}")print(response.text)return Noneif __name__ == "__main__":text_to_predict = "探索神秘的海底世界"prediction = predict(text_to_predict)if prediction:print(f"Prediction: {prediction['prediction']}")

详细说明

1. 数据处理模块：

   • 读取训练集、验证集和测试集的数据文件。
   • 使用BERT的Tokenizer对数据进行tokenize处理，生成模型可接受的输入格式。
   • 提供主要的数据处理函数，包括加载数据和tokenize数据。

2. 数据集类模块：

   • 定义一个TextDataset类，用于将tokenized数据和标签封装成PyTorch的数据集格式，便于Trainer进行训练和评估。

3. 模型训练模块：

   • 使用数据处理模块加载和tokenize数据。
   • 创建训练和验证数据集。
   • 加载DistilBERT模型，并设置训练参数（包括启用混合精度训练）。
   • 使用Trainer进行模型训练和评估，并保存训练好的模型。

4. 模型部署模块：

   • 使用FastAPI创建一个简单的API服务。
   • 加载保存的模型和tokenizer。
   • 定义一个预测接口，通过POST请求接收文本输入并返回分类预测结果。

5. 主脚本：

   • 运行模型训练流程，并输出训练完成的提示。