ONNX 模型 - 一个经过微调的 DistilBERT 版本，可用于将文本分类为以下类别之一：

• neutral（中立）、offensive_language（冒犯性语言）、harmful_behaviour（有害行为）、hate_speech（仇恨言论）

该模型使用 csfy 工具 和数据集 seanius/toxic-or-neutral-text-labelled 进行训练。

需要基础模型（distilbert-base-uncased）。

有关如何运行该模型的示例，请参见下方 - 或参阅 csfy 工具。

输出是一个表示类别的数字 - 通过 label_mapping.json 文件进行解码。

使用方法

# Loading the label mappings
import json
def load_label_mappings():
    with open("./label_mapping.json", encoding="utf-8") as f:
        data = json.load(f)
        return data['labels']

label_mappings = load_label_mappings()

# Loading the model
import onnxruntime as ort
from transformers import DistilBertTokenizer
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
ort_session = ort.InferenceSession("./toxic-or-neutral-text-labelled.onnx")

# Predicting label for given text
def predict_via_onnx(text, ort_session, tokenizer, label_mappings):
    model_expected_input_shape = ort_session.get_inputs()[0].shape
    print("Model expects input shape:", model_expected_input_shape)
    inputs = tokenizer(text, return_tensors="np", padding="max_length", truncation=True, max_length=model_expected_input_shape[1])
    print("input shape", inputs['input_ids'].shape)

    input_ids = inputs['input_ids']
    if input_ids.ndim == 1:
        input_ids = input_ids[np.newaxis, :]
    ort_inputs = {ort_session.get_inputs()[0].name: input_ids}

    ort_inputs['input_ids'] = ort_inputs['input_ids'].astype(np.int64)

    ort_outputs = ort_session.run(None, ort_inputs)
    predictions = np.argmax(ort_outputs, axis=-1)

    predicted_label = label_mappings[predictions.item()]
    return predicted_label

predicted_label = predict_via_onnx("How do I get to the beach?", ort_session, tokenizer, label_mappings)
print(predicted_label)

license: mit