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

🕗 发布于 2024-07-07 09:19 自然语言处理 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

文件说明

bert_data/：存放训练集、验证集和测试集的数据文件。
- train.txt
- dev.txt
- test.txt
saved_model/：存放训练好的模型和tokenizer。
results/：存放训练结果。
logs/：存放训练日志。
data_processing.py：数据处理模块，负责读取和预处理数据。
dataset.py：数据集类模块，定义了用于训练和评估的数据集类。
training.py：模型训练模块，定义了训练和评估模型的过程。
app.py：模型部署模块，使用FastAPI创建API服务。
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(): 运行模型训练流程，并输出训练完成的提示。

运行步骤

安装依赖

pip install pandas torch transformers fastapi uvicorn scikit-learn

数据处理

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

训练模型

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

python main.py

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

部署模型

运行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 DistilBertTokenizer

def load_data(file_path):
    data = pd.read_csv(file_path, delimiter='\t', header=None)
    data.columns = ['text', 'label']
    return data

def tokenize_data(data, tokenizer, max_length=128):
    encodings = tokenizer(list(data['text']), truncation=True, padding=True, max_length=max_length)
    return encodings

def main():
    # 加载Tokenizer
    tokenizer = 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 torch

class TextDataset(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)

3. 模型训练模块

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

# training.py
import torch
from transformers import DistilBertForSequenceClassification, Trainer, TrainingArguments
from dataset import TextDataset
import data_processing

def train_model():
    # 加载数据和tokenizer
    train_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,                       # 启用混合精度训练
    )

    # 定义Trainer
    trainer = 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.tokenizer
    tokenizer.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 torch

app = 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 uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

5. 主脚本

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

# main.py
import training

def main():
    # 训练模型
    training.train_model()
    print("模型训练完成并保存。")

if __name__ == "__main__":
    main()

其他：客户端调用案例

# client.py
import requests

def 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 prediction
    else:
        print(f"Error: {response.status_code}")
        print(response.text)
        return None

if __name__ == "__main__":
    text_to_predict = "探索神秘的海底世界"
    prediction = predict(text_to_predict)
    if prediction:
        print(f"Prediction: {prediction['prediction']}")

详细说明

数据处理模块：
- 读取训练集、验证集和测试集的数据文件。
- 使用BERT的Tokenizer对数据进行tokenize处理，生成模型可接受的输入格式。
- 提供主要的数据处理函数，包括加载数据和tokenize数据。
数据集类模块：
- 定义一个TextDataset类，用于将tokenized数据和标签封装成PyTorch的数据集格式，便于Trainer进行训练和评估。
模型训练模块：
- 使用数据处理模块加载和tokenize数据。
- 创建训练和验证数据集。
- 加载DistilBERT模型，并设置训练参数（包括启用混合精度训练）。
- 使用Trainer进行模型训练和评估，并保存训练好的模型。
模型部署模块：
- 使用FastAPI创建一个简单的API服务。
- 加载保存的模型和tokenizer。
- 定义一个预测接口，通过POST请求接收文本输入并返回分类预测结果。
主脚本：
- 运行模型训练流程，并输出训练完成的提示。

原文地址：https://blog.csdn.net/weixin_47552266/article/details/140161623

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！

上一篇：软考中级数据库系统工程师备考经验分享
下一篇：ESP32CAM物联网教学04

leetcode289:生命游戏
根据，简称为，是英国数学家约翰·何顿·康威在 1970 年发明的细胞自动机。给定一个包含m × n个格子的面板，每一个格子都可以看成是一个细胞。每个细胞都具有一个初始状态：1即为（live），或0即为
阅读更多2024-10-20
MongoDB数据恢复
注意：两个MongoDB的版本要一致，本文使用的是mongo:4.2.24。先把K8S上面的MongoDB 容器停止（可以把副本改成0）。1、将容器挂载MongoDB的数据目录备份到本地。经常是数据文
阅读更多2024-10-20
C#中实现事务
C#中实现事务
阅读更多2024-10-20
【LeetCode每日一题】——560.和为 K 的子数组
给你一个整数数组 nums 和一个整数 k ，请你统计并返回该数组中和为 k 的子数组的个数。子数组是数组中元素的连续非空序列。
阅读更多2024-10-20
「漏洞复现」满客宝智慧食堂系统 selectUserByOrgId 未授权访问漏洞
请勿利用文章内的相关技术从事非法测试，由于传播、利用此文所提供的信息而造成的任何直接或者间接的后果及损失，均由使用者本人负责，作者不为此承担任何责任。工具来自网络，安全性自测，如有侵权请联系删除。本次
阅读更多2024-10-20
React面试题目（从基本到高级）
React前端面试常见题目涵盖了React的基础概念、组件、状态管理、生命周期、性能优化等多个方面。
阅读更多2024-10-20
12.个人博客系统（Java项目基于spring和vue）
1 在校学习的学生，可用于日常学习使用或是毕业设计使用 2 毕业一到两年的开发人员，用于锻炼自己的独立功能模块设计能力，增强代码编写能力。 3 亦可以部署为商化项目使用。 4 需要完整资料及源码
阅读更多2024-10-20
YoloV8改进策略：注意力改进|DeBiFormer，可变形双级路由注意力|引入DeBiLevelRoutingAttention注意力模块（全网首发）
本次改进的核心在于将DeBiLevelRoutingAttention模块嵌入到YoloV8的主干网络中，具体位置是在SPPF（Spatial Pyramid Pooling Fast）模块之后。这一
阅读更多2024-10-20
word取消自动单词首字母大写
情况说明：在word输入单词后首字母会自动变成大写取消单词首字母大写步骤：（1）点击菜单栏文件（2）点击“更多”——>“选项”（3）点击“校对”——>“自动更正选项”（4）取消“句首字母大
阅读更多2024-10-20
web前端网页用户注册页面
【代码】web前端网页用户注册页面。
阅读更多2024-10-20

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

项目说明

项目名称

项目概述

项目结构

文件说明

数据集数据规范

数据文件格式

数据示例

注意事项

模块说明

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

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

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

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

5. 主脚本 (main.py)

运行步骤

示例请求

注意事项

项目代码

1. 数据处理模块

2. 数据集类模块

3. 模型训练模块

4. 模型部署模块

5. 主脚本

其他：客户端调用案例

详细说明

相关文章