时间戳

首页 > 解决方案 > 使用 Huggingface TextClassificationPipeline 时如何设置标签名称？

nlp - 使用 Huggingface TextClassificationPipeline 时如何设置标签名称？

问题描述

我正在使用带有 TextClassificationPipeline 的微调 Huggingface 模型（在我的公司数据上）来进行类预测。现在，此Pipeline预测的标签默认为LABEL_0，LABEL_1依此类推。有没有办法为TextClassificationPipeline对象提供标签映射，以便输出反映相同？

环境：

  • 张量流==2.3.1
  • 变压器==4.3.2

示例代码：

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # or any {'0', '1', '2'}

from transformers import TextClassificationPipeline, TFAutoModelForSequenceClassification, AutoTokenizer

MODEL_DIR = "path\to\my\fine-tuned\model"

# Feature extraction pipeline
model = TFAutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)

pipeline = TextClassificationPipeline(model=model,
                                      tokenizer=tokenizer,
                                      framework='tf',
                                      device=0)

result = pipeline("It was a good watch. But a little boring.")[0]

输出：

In [2]: result
Out[2]: {'label': 'LABEL_1', 'score': 0.8864616751670837}

标签： nlphuggingface-transformers

解决方案

添加此类映射最简单的方法是编辑模型的 config.json 以包含：id2label字段，如下所示：

{
  "_name_or_path": "distilbert-base-uncased",
  "activation": "gelu",
  "architectures": [
    "DistilBertForMaskedLM"
  ],
  "id2label": [
    "negative",
    "positive"
  ],
  "attention_dropout": 0.1,
  .
  .
}

设置此映射的代码内方法是id2label在调用中添加参数，from_pretrained如下所示：

model = TFAutoModelForSequenceClassification.from_pretrained(MODEL_DIR, id2label={0: 'negative', 1: 'positive'})

这是我为此提出的Github 问题，以添加到 transformers.XForSequenceClassification 的文档中。

推荐阅读