pronics2004's picture
Update README.md
3ed34fd verified
|
raw
history blame
1.99 kB
metadata
license: apache-2.0
language:
  - en
pipeline_tag: text-classification

Model Description

This model is IBM's lightweight, 4-layer toxicity binary classifier for English. Its latency characteristics make it a suitable guardrail for any large language model. It can also be used for bulk processing of data where high throughput is needed. It has been trained on several benchmark datasets in English, specifically for detecting hateful, abusive, profane and other toxic content in plain text.

Feature

This model offers very low inference latency and is capable of running on CPUs apart from GPUs and AIUs. It features 38 million parameters, reducing the number of hidden layers from 12 to 4, decreasing the hidden size from 768 to 576, and the intermediate size from 3072 to 768, compared to the original RoBERTa model architecture.

Description of Image

Model Usage

# Example of how to use the model
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name_or_path = 'ibm-granite/granite-guardian-hap-38m'
model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

# Sample text
text = ["This is the 1st test", "This is the 2nd test"]
input = tokenizer(text, padding=True, truncation=True, return_tensors="pt")

with torch.no_grad():
    logits = model(**input).logits
    prediction = torch.argmax(logits, dim=1).detach().numpy().tolist() # Binary prediction where label 1 indicates toxicity.
    probability = torch.softmax(logits, dim=1).detach().numpy()[:,1].tolist() #  Probability of toxicity.
   

Performance Comparison with Other Models

The model delivers competitive performance with significantly lower inference latency. If a better F1 score is required, please refer to IBM's 12-layer model here.

Description of Image Description of Image