hate-speech-and-offensive-message-classifier

import re
import html
import contractions
from transformers import RobertaTokenizer, RobertaForSequenceClassification
import torch

# Load the trained model + tokenizer
model = RobertaForSequenceClassification.from_pretrained("AshiniR/hate-speech-and-offensive-message-classifier")
tokenizer = RobertaTokenizer.from_pretrained("AshiniR/hate-speech-and-offensive-message-classifier")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

def preprocess_text(text: str) -> str:
    """
    Preprocess raw text for transformer-based models like RoBERTa.

    This function is tailored for toxicity, sentiment, and social media classification.
    It removes noise (URLs, mentions, HTML codes) but keeps important signals
    such as casing, punctuation, and emojis.

    Steps:
        1. Decode HTML entities (e.g., '>' → '>')
        2. Replace URLs with placeholders ("")
        3. Replace mentions with placeholders ("")
        4. Remove '#' from hashtags but keep the word (e.g., "#love" → "love")
        5. Expand contractions (e.g., "you're" → "you are")
        6. Mildly normalize repeated characters (3+ → 2)
        7. Remove "RT" only if at start of tweet
        8. Normalize whitespace

    Args:
        text (str): Raw tweet text.

    Returns:
        str: Cleaned text suitable for RoBERTa tokenization.
    """
    if not isinstance(text, str):
        return ""

    # 1. Decode HTML entities
    text = html.unescape(text)

    # 2. Replace URLs with placeholder
    text = re.sub(r"(https?://\S+|www\.\S+)", "", text)

    # 3. Replace user mentions with placeholder
    text = re.sub(r"@\w+", "", text)

    # 4. Simplify hashtags
    text = re.sub(r"#(\w+)", r"\1", text)

    # 5. Expand contractions
    text = contractions.fix(text)

    # 6. Mild normalization of character elongations (3+ → 2)
    text = re.sub(r"(.)\1{2,}", r"\1\1", text)

    # 7. Remove RT only if it starts the tweet (For tweets)
    text = re.sub(
        r"^[\s\W]*rt\s*@?\w*:?[\s-]*",
        "",
        text,
        flags=re.IGNORECASE
    )

    # 8. Normalize whitespace
    text = re.sub(r"\s+", " ", text).strip()

    return text


def get_inference(text: str) -> list:
    """Returns prediction results in [{'label': str, 'score': float}, ...] format."""
    # Preprocess the text
    text = preprocess_text(text)

    # Tokenize input text
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=False,
        max_length=128
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    # Get model predictions
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=-1)
    
    # Convert to label format
    labels = ["neither", "hate/offensive"]
    results = []
    for i, prob in enumerate(probabilities[0]):
        results.append({
            "label": labels[i],
            "score": prob.item()
        })
    
    return sorted(results, key=lambda x: x["score"], reverse=True)

# Example usage
text = "your example massege"
predictions = get_inference(text)
print(f"Text: '{text}'")
print(f"Predictions: {predictions}")