Build-IT-of-AZ/TinyBERT-Multi-Head-Merged-Data-Classifier-onnx-v1

TinyBERT Multi-Head Binary Classifier ONNX

This is the ONNX version of the TinyBERT Multi-Head Binary Classifier for content moderation.

Original Model

• Original Repository: Build-IT-of-AZ/TinyBERT-Multi-Head-Merged-Data-Classifier-v1
• Base Model: huawei-noah/TinyBERT_General_4L_312D
• Architecture: Multi-head binary classifier with 9 heads

Classes

The model can detect the following content types:

• Normal (default when no other class is detected)
• Vulgarity
• Hate Speech
• Bullying
• Self Harm
• Substance Abuse
• PII (Personally Identifiable Information)
• Platform Departure
• Mental Health
• Grooming

Usage

import onnxruntime as ort
from transformers import AutoTokenizer
import numpy as np

# Load the ONNX model
session = ort.InferenceSession("model.onnx")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("huawei-noah/TinyBERT_General_4L_312D")

# Tokenize input
text = "Your input text here"
inputs = tokenizer(text, truncation=True, padding='max_length', max_length=512, return_tensors='pt')

# Convert to numpy for ONNX
input_feed = {name: np.array(value) for name, value in inputs.items()}

# Run inference
outputs = session.run(None, input_feed)
logits = outputs[0]

# Apply sigmoid to get probabilities
probabilities = 1 / (1 + np.exp(-logits))

# Apply threshold (default 0.5) to get binary predictions
binary_preds = probabilities > 0.5

Model Details

• Input: Tokenized text (max length: 512)
• Output: 9 logits corresponding to each non-normal class
• Threshold: 0.5 (configurable)
• Framework: ONNX Runtime compatible