情绪分类模型训练项目
February 26, 2026 · View on GitHub
本项目用于训练一个能够识别19+种不同情绪的文本分类模型。
项目结构
New-Emotion-Trainer/
├── data/ # 数据目录
│ └── emotion_data_manual.csv
├── pre-process/ # 数据预处理
│ └── csvCleaner.py
├── train/ # 训练模块
│ ├── __init__.py
│ ├── config.py # 配置模块
│ ├── data_loader.py # 数据加载模块
│ ├── dataset.py # 数据集类
│ └── trainer.py # 训练模块
├── inference/ # 推理模块
│ ├── __init__.py
│ ├── config.py # 推理设置模块
│ └── emotion_predictor.py# 推理类模块
├── .env # 环境变量配置
├── main.py # 主程序入口
├── requirements.txt # 依赖列表
└── README.md # 项目说明
安装依赖
pip install -r requirements.txt
配置
项目使用.env文件进行配置,包含以下主要配置项:
- 数据路径
- 模型参数
- 训练超参数
- 输出路径
- 情绪标签列表
- NEFTune噪声强度 (NEFTUNE_NOISE_ALPHA): 设置为0.0表示禁用NEFTune
使用方法
运行训练程序:
python main.py
情绪类别
模型可以识别以下19种情绪: 可以在env和data中的csv新增识别情绪
- 高兴
- 厌恶
- 害羞
- 害怕
- 生气
- 认真
- 紧张
- 慌张
- 疑惑
- 兴奋
- 无奈
- 担心
- 惊讶
- 哭泣
- 心动
- 难为情
- 自信
- 调皮
- 平静
NEFTune支持
本项目支持NEFTune技术,通过在嵌入层添加噪声来提高模型性能。要启用NEFTune,请在.env文件中设置NEFTUNE_NOISE_ALPHA为正值(如0.1)。注意:需要transformers>=4.33.0版本。