文本安全评分
更新时间:2026-06-15
简介
文本安全性评分器 - 基于ShieldLM-6B-chatglm3的安全性评估
功能描述
- 多语言支持**:支持中文和英文文本安全性评估
- 三分类评估**:输出safe、unsafe、controversial三类概率
- 批量处理**:支持批量文本安全性评估,提升处理效率
- 模型核心**:基于ShieldLM-6B-chatglm3
- 推理优化**:支持GPU加速和批量推理
- 内容安全审核
- 文本风险评估
- 多语言安全过滤
算子参数
输入
| 输入 | 含义 |
|---|---|
| texts | pyarrow.Array,元素类型为str |
输出
| 输出 | 含义 |
|---|---|
| safe | Float64,模型预测文本为“安全”的概率 |
| unsafe | Float64,模型预测文本为“不安全”的概率 |
| controversial | Float64,模型预测文本为“有争议”的概率 |
参数
| 参数名称 | 类型 | 默认值 | 描述 |
|---|---|---|---|
| lang | str | 'zh' | 语种 描述:需要评估的文本的语种 可选值:["en", "zh"] 默认值:"zh" |
| model_path | str | '/opt/aihc/models' | 模型文件所在的基础路径 默认值:"/opt/aihc/models" |
| model_name | str | 'thu-coai/ShieldLM-6B-chatglm3' | 模型名称 默认值:"thu-coai/ShieldLM-6B-chatglm3" |
| batch_size | int | 1 | 批处理大小 描述:控制模型推理时的批量处理大小 影响:较大的batch_size可以提高GPU利用率和吞吐量,但会增加显存占用 建议:ShieldLM-6B模型显存占用较大,默认设置为1以确保稳定性和兼容性 调优:在显存充足的环境下可适当增加到2-4以提升处理效率 默认值:1 |
| rank | int | None | GPU编号 描述:指定使用的GPU编号,None表示自动选择 默认值:None |
调用示例
Python
1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.text.text_safety_scorer import TextSafetyScorer
10
11if __name__ == "__main__":
12 if os.getenv("DAFT_RUNNER", "native") == "ray":
13 import ray
14 ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15 daft.set_runner_ray()
16 daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18 # TODO: 根据实际场景准备样本数据
19 samples = {"texts": [...]}
20 ds = daft.from_pydict(samples)
21 constructor_kwargs = {
22 "lang": 'zh',
23 "model_path": '/opt/aihc/models',
24 "model_name": 'thu-coai/ShieldLM-6B-chatglm3',
25 "batch_size": 1,
26 "rank": None,
27 }
28 ds = ds.with_column(
29 "result",
30 aihc_udf(
31 TextSafetyScorer,
32 construct_args=constructor_kwargs,
33 num_cpus=1,
34 concurrency=4,
35 batch_size=8,
36 )(col("texts")),
37 )
38 ds.show()
评价此篇文章
