项目说明

项目名称

基于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 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']}")

详细说明

1. 数据处理模块：

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

2. 数据集类模块：

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

3. 模型训练模块：

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

4. 模型部署模块：

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

5. 主脚本：

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