# nlp-classification

**Repository Path**: xtoll/nlp-classification

## Basic Information

- **Project Name**: nlp-classification
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-01-21
- **Last Updated**: 2026-01-21

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 工业电气设备文本分类系统

基于 BERT 的中文工业电气设备短文本分类项目，用于自动识别和分类低压电气设备类型。

## 📋 项目简介

本项目使用 `hfl/chinese-roberta-wwm-ext` 预训练模型进行微调，实现对工业电气设备名称的自动分类。支持多种低压电气设备类型识别，如断路器、接触器、电容器、互感器等。

**核心特性：**
- 🚀 基于中文 RoBERTa 预训练模型，专门优化中文理解能力
- 📊 自动处理类别不平衡问题（加权 F1 评估 + 分层采样）
- 💾 内置最佳模型自动保存机制
- 🔥 支持混合精度训练（FP16），加速训练过程
- 📈 实时评估指标监控（准确率 + F1 分数）

## 🗂️ 项目结构

```
nlp-classification/
├── train_ner.py                 # 模型训练脚本
├── predict.py                   # 批量预测脚本
├── pyproject.toml               # 项目依赖配置
├── datasets/
│   ├── train/                   # 训练数据目录
│   │   └── project-3-*.json     # JSON 格式训练数据
│   └── test_data.txt            # 测试数据（每行一个样本）
├── bert_industrial_model/       # 训练输出目录
│   ├── checkpoint-*/            # 训练中间检查点
│   └── final_model/             # 最终模型文件
│       ├── config.json
│       ├── model.safetensors
│       ├── tokenizer_config.json
│       └── vocab.txt
└── README.md                    # 本文档
```

## 🔧 环境配置

### 系统要求
- Python >= 3.10
- CUDA（可选，用于 GPU 加速）

### 安装依赖

使用 `uv`（推荐）：
```bash
uv sync
```

或使用 `pip`：
```bash
pip install -r requirements.txt
```

主要依赖包：
- `transformers >= 4.57.3` - HuggingFace 核心库
- `torch >= 2.9.1` - PyTorch 深度学习框架
- `datasets >= 4.4.2` - 数据集处理
- `evaluate >= 0.4.6` - 模型评估指标
- `scikit-learn >= 1.7.2` - 数据划分与评估
- `pandas` - 数据处理

## 📊 数据格式

### 训练数据格式（JSON）
```json
[
  {
    "text": "低压电流互感器",
    "sentiment": "低压电流互感器"
  },
  {
    "text": "塑壳式断路器",
    "sentiment": "塑壳式断路器"
  }
]
```

**字段说明：**
- `text`: 输入文本（设备名称或描述）
- `sentiment`: 类别标签（与文本对应的设备类型）

### 测试数据格式（TXT）
```text
低压电流互感器
多功能表
塑壳式断路器
```
每行一个待预测的文本样本。

## 🚀 快速开始

### 1. 训练模型

修改 `train_ner.py` 中的配置参数：

```python
CONFIG = {
    "data_file": "./datasets/train/project-3-at-2025-12-30-14-02-726444b7.json",
    "model_checkpoint": "hfl/chinese-roberta-wwm-ext",
    "output_dir": "./bert_industrial_model",
    "max_len": 64,          # 最大序列长度
    "batch_size": 16,       # 批次大小（显存不足可改为 8）
    "epochs": 50,           # 训练轮数
    "learning_rate": 3e-5,  # 学习率
}
```

执行训练：
```bash
python train_ner.py
```

**训练过程：**
1. 自动清洗数据并过滤稀有类别（样本数 < 2）
2. 按 8:2 比例划分训练集和验证集（分层采样）
3. 每个 epoch 结束后评估模型性能
4. 自动保存最佳模型到 `./bert_industrial_model/final_model/`

### 2. 模型预测

准备测试文件 `datasets/test_data.txt`，执行：
```bash
python predict.py
```

**输出示例：**
```
文本: 低压电流互感器        -> 预测分类: 低压电流互感器 (置信度: 0.9987)
文本: 塑壳 DsfC器          -> 预测分类: 塑壳式断路器 (置信度: 0.8923)
文本: 多功能表             -> 预测分类: 多功能表 (置信度: 0.9956)
```

## 📈 模型评估

训练过程中自动计算以下指标：

| 指标 | 说明 |
|------|------|
| **Accuracy** | 准确率 - 预测正确的样本占总样本的比例 |
| **F1-Score (Weighted)** | 加权 F1 分数 - 适用于类别不平衡场景 |

评估策略：
- 每个 epoch 结束后评估一次
- 以准确率（`accuracy`）作为最佳模型选择标准
- 仅保留效果最好的一个模型检查点

## ⚙️ 核心功能说明

### 数据处理流程
1. **数据清洗**：去除空值、重命名列
2. **稀有类别过滤**：剔除样本数 < 2 的类别
3. **标签映射**：自动构建 `label2id` 和 `id2label` 字典
4. **分层采样**：确保训练集和验证集类别分布一致

### 模型架构
- **基座模型**：`hfl/chinese-roberta-wwm-ext`（中文全词掩码预训练）
- **任务头**：序列分类头（线性层 + Softmax）
- **优化器**：AdamW（权重衰减 0.01）
- **学习率调度**：线性衰减

### 训练优化
- **动态补齐**：`DataCollatorWithPadding` 节省显存
- **混合精度训练**：自动检测 GPU 并启用 FP16
- **梯度累积**：支持小显存设备训练大模型
- **早停机制**：自动加载最佳模型权重

## 🔍 支持的设备类型

当前数据集支持的电气设备分类包括：
- 低压电流互感器
- 多功能表
- 电涌后备保护断路器
- 电涌保护器
- 微型断路器
- 塑壳式断路器
- 低压交流接触器
- 电容器投切开关
- 低压并联电容器
- 刀开关
- ...更多类型见训练数据

## 🛠️ 常见问题

**Q: 显存不足怎么办？**  
A: 减小 `batch_size`（8 → 4）或 `max_len`（64 → 32）

**Q: 训练速度慢？**  
A: 确认已安装 CUDA 版本 PyTorch，启用 FP16 混合精度训练

**Q: 如何添加新类别？**  
A: 在训练数据中添加新类别样本（建议 ≥ 10 个），重新训练即可

**Q: 预测结果置信度低？**  
A: 增加训练数据量、调整学习率或增加训练轮数

## 📝 开发说明

### 自定义配置
在 `train_ner.py` 的 `CONFIG` 字典中修改参数：
- `max_len`: 根据文本长度调整（建议 32-128）
- `epochs`: 观察验证集准确率，防止过拟合
- `learning_rate`: 模型不收敛可尝试 2e-5 或 5e-5

### 代码结构
```python
# 1. 配置中心
CONFIG = {...}

# 2. 辅助函数
compute_metrics()      # 评估指标计算
preprocess_function()  # 数据预处理

# 3. 主程序
main()                 # 训练流程编排
```

## 📄 许可证

本项目仅供学习和研究使用。

## 🤝 贡献

欢迎提交 Issue 和 Pull Request！

---

**最后更新时间**：2026年1月21日