# Ours_TMF

**Repository Path**: dxatar/ours-ai

## Basic Information

- **Project Name**: Ours_TMF
- **Description**: 小组项目,合力完成,力争第一
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-09-06
- **Last Updated**: 2026-09-12

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 投满分 · 今日头条短文本多分类项目

> 面向今日头条推荐场景的短文本自动分类项目（把新闻 / 资讯短文本自动分到对应频道，再“投递”进各自的推荐流）。

---

## 一、项目背景

字节跳动旗下产品线众多，其中**抖音**与**今日头条**是最大的两张王牌，分别代表"基于短视频的推荐"与"基于短文本的推荐"，背后融合了 CV、NLP、推荐系统、大数据等多个领域的技术。以今日头条为例：用户在大量资讯中，必然对某些类别更感兴趣（如男生的历史、军事、足球;女生的财经、八卦、美妆）。如果能主动把用户更感兴趣的类别筛选出来并做推荐阅读，点击量、订阅量、付费量都会明显增长。因此，在今日头条 App 的推荐系统模块中，需要内嵌一个子任务：**将短文本自动多分类，并像"投递"一样把内容送到对应的"频道 / 推荐流"中**。"投满分"项目正是这一功能子集——它只负责完成新闻、资讯等短文本的多分类。

## 二、功能概述

本项目采用经典的 **离线训练 + 在线推理** 分离架构：

- **离线侧**：对原始语料做预处理，微调 BERT 分类模型，并通过**知识蒸馏**压缩为小模型。
- **在线侧**：以 Flask 提供 HTTP 服务，支持**单条输入 / 批量输入 / 文件上传**三种方式获取分类结果。
- **前端**：Vue3 + Vite 单页应用，由 Nginx 托管并反向代理接口。
- **部署**：Docker Compose 双容器编排（API 容器 + Web 容器），一条命令启动。

| 环节 | 入口 | 说明 |
| --- | --- | --- |
| 离线训练流水线 | `run_pipeline.py` | 编排 数据准备 → 训练导出 → 模型评估 → 蒸馏 → 蒸馏评估 |
| 在线推理服务 | `wsgi.py` | Flask 应用启动入口（应用工厂） |
| 分类接口 | `POST /api/v4/classify` | 短文本 → 类别 + 置信度 + 候选分布 |
| 批量接口 | `POST /api/v4/classify/batch` | 多行文本 → 逐条结果 + 类别分布 |
| 文件接口 | `POST /api/v4/classify/file` | 上传 txt / csv → 批量结果（带标签时统计准确率） |
| 前端页面 | `front/` | 单条 / 批量 / 上传三种交互 |

## 三、技术栈

- **语言**：Python 3.12（开发环境为 conda 环境 `chengdu_q3_tmf`）、JavaScript（前端）
- **建模**：PyTorch（BERT 微调、知识蒸馏、动态量化、结构化剪枝）、Transformers、ModelScope（预训练权重）
- **在线服务**：Flask + flask-cors（应用工厂 + 蓝图 + 版本化路由）、Gunicorn
- **前端**：Vue 3 + Vite + axios
- **部署**：Docker / Docker Compose + Nginx

### 模型演进路线

| 阶段 | 产物 | 说明 |
| --- | --- | --- |
| 微调 | `model/ft_bert_model/` | Chinese-MacBERT 微调，12 层，冻结前 6 层 |
| 蒸馏 | `model/ft_bert_distilled/` | 以微调模型为教师、rbt3 为学生，3 层，**在线服务采用此模型** |
| 量化 / 剪枝 | `model/ft_bert_quantized/`、`model/ft_bert_pruned/` | 动态量化与结构化剪枝（实验性，见 `src/utils.py`） |

> 在线推理选用蒸馏后的 3 层模型：实测单条推理约 **6.9ms**，批大小 64 时吞吐约 **344 条/秒**，非常适合 CPU 部署。

## 四、数据集

数据为今日头条渠道的新闻标题（短文本）语料，位于 `data/raw/` 下，`\t` 分隔，一行一条：

```
标题文本\t类别编号
```

| 文件 | 行数 | 用途 |
| --- | --- | --- |
| `data/raw/tmf_train.txt` | 180,000 | 训练集 |
| `data/raw/tmf_valid.txt` | 10,000 | 验证集 |
| `data/raw/tmf_test.txt` | 10,000 | 测试集 |

- 类别共 10 个，清单见 `data/tmf_class.txt`，**顺序即标签编号（0-based）**：
  `finance`(0) / `realty`(1) / `stocks`(2) / `education`(3) / `science`(4) / `society`(5) / `politics`(6) / `sports`(7) / `game`(8) / `entertainment`(9)。
- 注意：模型 `config.json` 中的 `id2label` 是未赋名的 `LABEL_0..LABEL_9`，真实类名一律以 `tmf_class.txt` 为准。
- `data/stopwords.txt`、`data/user_dict.txt` 为早期分词方案（jieba）的遗留文件，当前 BERT 方案未使用。

## 五、目录结构

```
ours-ai/
├── run_pipeline.py         # 离线训练流水线入口
├── wsgi.py                 # Web 服务启动入口（gunicorn 挂载 wsgi:app）
├── requirements.txt        # 完整依赖（离线训练 + 在线服务）
├── requirements-serve.txt  # 在线服务专属依赖（用于构建精简镜像）
├── Dockerfile              # API 容器镜像
├── docker-compose.yml      # 双容器编排（api + web）
├── .env.example            # 环境变量模板
├── src/                    # —— 离线模型侧 ——
│   ├── config.py           # 训练侧配置（路径、超参，dataclass）
│   ├── model_train.py      # BERT 微调训练与导出
│   ├── model_eval.py       # 模型评估（含 ONNX 导出示例）
│   └── utils.py            # 通用工具 + 蒸馏 / 量化 / 剪枝
├── app/                    # —— 在线服务侧（Flask）——
│   ├── __init__.py         # 应用工厂 create_app()：配置/日志/CORS/蓝图/错误处理
│   ├── config.py           # Web 配置（dev / prod / test 环境映射）
│   ├── inference.py        # 推理引擎（进程内单例，设备自动选择）
│   ├── auth.py             # API Key 鉴权装饰器（默认关闭）
│   ├── parsers.py          # 上传文件解析（编码与行格式兼容）
│   ├── response.py         # 统一响应信封
│   └── v4/
│       └── prediction.py   # v4 分类接口蓝图（url_prefix=/api/v4）
├── front/                  # —— 前端（Vue3 + Vite）——
│   ├── src/api/            # axios 客户端与接口封装
│   ├── src/components/     # 三个输入面板 + 结果展示组件
│   └── Dockerfile          # 多阶段构建 → nginx 托管
├── nginx/
│   └── default.conf        # 静态资源托管 + /api 反向代理
├── data/
│   ├── raw/                # 原始语料（训练/验证/测试）
│   ├── pre/                # 早期 fastText 预处理产物（当前未使用）
│   ├── stopwords.txt       # 停用词（遗留）
│   └── tmf_class.txt       # 类别清单（顺序即标签编号）
├── model/                  # 训练产物（已 gitignore，需本地生成或另行分发）
├── logs/                   # 日志目录
└── tests/                  # 接口测试（pytest）
```

## 六、安装与运行

```bash
# 1. 准备 Python 3.12 环境
conda create -n chengdu_q3_tmf python=3.12
conda activate chengdu_q3_tmf

# 2. 安装依赖
pip install -r requirements.txt
```

> 模型文件位于 `model/` 且被 `.gitignore` 忽略，克隆仓库后需自行训练生成，或从交付介质中拷贝。

## 七、使用说明

### 1. 离线训练流水线

```bash
python run_pipeline.py
```

流水线规划了五个步骤（对应 `run_pipeline.py`）：

1. **数据清洗与准备**（规划中，`src/data_pre.py` 尚未实现，当前为注释状态）;
2. **训练并导出分类模型**（`src/model_train.py` → `model/ft_bert_model/`）;
3. **模型评估**（`src/model_eval.py`）;
4. **模型压缩：知识蒸馏**（`src/utils.py::distill_model` → `model/ft_bert_distilled/`）;
5. **蒸馏模型评估**（`src/utils.py::evaluate_distilled_model`）。

> 训练与蒸馏耗时较长，按需取消 `run_pipeline.py` 中对应步骤的注释；已训练好的模型可直接跳过 2~4 步。

### 2. 启动在线服务（开发模式）

```bash
python wsgi.py          # 默认 127.0.0.1:8000
```

接口调用示例：

```bash
# 单条分类
curl -X POST http://127.0.0.1:8000/api/v4/classify \
     -H 'Content-Type: application/json' \
     -d '{"text":"CBA总决赛辽宁队夺冠 郭艾伦砍下30分"}'

# 批量分类
curl -X POST http://127.0.0.1:8000/api/v4/classify/batch \
     -H 'Content-Type: application/json' \
     -d '{"texts":["央行宣布降准0.5个百分点","新款游戏上线首日登顶"],"top_k":3}'

# 文件上传（txt / csv，每行一条，可带标签）
curl -X POST http://127.0.0.1:8000/api/v4/classify/file \
     -F 'file=@data/raw/tmf_test.txt' -F 'top_k=3'

# 健康检查
curl http://127.0.0.1:8000/api/v4/health
```

**生产部署参考**（务必置于 Nginx 等反向代理之后）：

```bash
gunicorn -w 4 -b 0.0.0.0:8000 --preload -t 300 wsgi:app
```

> `--preload` 让模型只在 master 进程加载一次，fork 后各 worker 通过写时复制共享内存，4 个 worker 实测总占用约 330MB 而非叠加。

### 3. 前端开发

```bash
cd front
npm install
npm run dev     # http://localhost:5173，/api 会自动代理到 127.0.0.1:8000
```

### 4. Docker 一键部署（推荐）

```bash
cp .env.example .env      # 按需修改配置
docker compose up --build -d
# 访问 http://localhost:8080
```

容器编排说明：

- `api` 容器：Gunicorn + Flask + 蒸馏模型，**不对外发布端口**，只在内网提供服务;
- `web` 容器：Nginx 托管前端静态页，并把 `/api` 反向代理到 `api:8000`;
- `web` 通过 `depends_on: service_healthy` 等待 `api` 健康检查通过后再启动。

## 八、接口说明

所有接口返回统一信封：`{"code": int, "message": str, "data": any}`，`code` 同时作为 HTTP 状态码。

| 状态码 | 含义 |
| --- | --- |
| 200 | 成功 |
| 400 | 参数错误（JSON 非法、`text` 缺失或为空、超长、`top_k` 越界） |
| 401 | 鉴权失败（仅在开启 API Key 时） |
| 413 | 上传内容超过大小限制（默认 10MB） |
| 422 | 无可处理内容（批量空数组 / 超批量上限 / 文件无有效行） |
| 500 | 服务器内部错误 |

**单条分类** `POST /api/v4/classify`

```json
// 请求
{ "text": "2009年成人高考招生统一考试时间表", "top_k": 5 }
// 响应 data
{
  "text": "2009年成人高考招生统一考试时间表",
  "label_id": 3,
  "label": "education",
  "confidence": 0.9731,
  "top_k": [{ "label_id": 3, "label": "education", "probability": 0.9731 }]
}
```

**文件上传** `POST /api/v4/classify/file`（`multipart/form-data`，字段 `file`，可选 `top_k`）

支持三种行格式，自动识别：

```
纯标题                      中华女子学院本科层次仅1专业招男生
标题 + Tab + 类别编号        中华女子学院本科层次仅1专业招男生	3
fastText（已分词）           __label__3 中华 女子 学院 本科 层次 仅 1 专业 招 男生
```

- 编码自动识别：优先 UTF-8（含 BOM），失败回退 GB18030;
- 首行表头（如 `标题\t类别编号`）自动跳过; 空行计入 `skipped`;
- 行内携带标签时，响应额外返回 `labeled` 与 `accuracy`（准确率），逐条结果含 `gold_label_id`。

## 九、配置说明

配置全部通过环境变量注入，模板见 `.env.example`。

- **训练侧配置**：`src/config.py` 中的 `Config`（路径、超参）。
- **服务侧配置**：`app/config.py` 中的 `config_map`，按 `FLASK_ENV` 区分 `development` / `production` / `test`。

常用环境变量：

| 变量 | 默认值 | 说明 |
| --- | --- | --- |
| `MODEL_DIR` | `model/ft_bert_distilled` | 在线推理使用的模型目录 |
| `MAX_LEN` | `32` | 序列截断长度 |
| `BATCH_SIZE` | `64` | 推理批大小 |
| `DEVICE` | `auto` | `auto` / `cuda` / `mps` / `cpu` |
| `MAX_CONTENT_LENGTH` | `10485760` | 上传大小上限（字节） |
| `MAX_BATCH_SIZE` | `500` | 单次批量上限（条） |
| `API_KEY_ENABLED` | `false` | 是否开启接口鉴权 |
| `API_KEY` | 空 | 开启鉴权后校验的密钥 |
| `LAZY_MODEL_LOAD` | `false` | 为 true 时首次请求才加载模型 |

> `API_KEY_ENABLED=true` 后，分类接口需携带请求头 `X-API-Key`（健康检查除外，以免影响容器探活）。前端可在「设置」中填入该密钥，保存在浏览器本地。

## 十、测试

```bash
pytest tests/ -v
```

测试通过 `tests/conftest.py` 中的 `FakePredictor` 替换真实推理引擎，**不会加载模型权重**，全部用例在 1 秒内完成。覆盖：三个接口的正常路径、各类错误码、编码兼容（UTF-8/BOM/GBK）、三种行格式、表头跳过、鉴权开关、统一 JSON 错误响应。

## 十一、实现要点

- **模型只加载一次**：`app/inference.py` 用模块级单例 + 双重检查锁，配合 Gunicorn `--preload` 实现多 worker 共享内存。
- **必须显式指定 `max_length`**：该模型 `tokenizer_config.json` 中 `model_max_length` 为未设置的极大值，不显式传参则 `truncation=True` 形同虚设。
- **离线可用**：模型以 `local_files_only=True` 加载；蒸馏模型目录缺失的 `vocab.txt` 在构建镜像时从基座补齐，确保断网环境可启动。
- **镜像精简**：在线镜像只安装 `requirements-serve.txt`（不含 modelscope / scikit-learn 等训练依赖），并移除 torch 自带的测试与头文件，镜像由 2.4GB 降至约 1.83GB。

---

Author： 高攀（PANGAO）
Date：2026-09-06