AI guardrails: Toxicity scorers

git clone https://github.com/wandb/weave.git && cd weave && git fetch origin pull/3006/head:xtra-scorers && git checkout xtra-scorers && pip install -qq -e . && pip install openai asyncio

from weave.scorers.moderation_scorer import ToxicityScorer
weave.init("toxicity-scorer")
toxicity_scorer = ToxicityScorer()

result = toxicity_scorer.score("This is a hateful message.")
print(result)
# {'flagged': False, 'categories': {'Race/Origin': 1, 'Gender/Sex': 0, 'Religion': 0, 'Ability': 0, 'Violence': 1}}
# sum < 5, categories < 2


result = toxicity_scorer.score("This is another hateful message.")
print(result)
# {'flagged': True, 'categories': {'Race/Origin': 2, 'Gender/Sex': 0, 'Religion': 0, 'Ability': 0, 'Violence': 1}}
# one category is 2>=

result = toxicity_scorer.score("This is a broad hateful message.")
print(result)
# {'flagged': True, 'categories': {'Race/Origin': 1, 'Gender/Sex': 1, 'Religion': 1, 'Ability': 1, 'Violence': 1}}
# sum >= 5


from openai import OpenAI
import weave; weave.init('toxicity-scorers')
client = OpenAI()


response = client.moderations.create(
  model="omni-moderation-latest",
  input="text to check toxicity for",
)

print(response)

from weave.scorers.llamaguard_scorer import LlamaGuardScorer
import asyncio

import weave; weave.init('toxicity-scorers')


async def main():
    # Initialize the LlamaGuardScorer with the 1B model
    scorer = LlamaGuardScorer(
        model_name_or_path="meta-llama/Llama-Guard-3-1B",
        device="cpu",  # Use "cuda" if a GPU is available
    )

    # Text to score
    sample_text = "Your input text here. Check if this is safe or not."

    # Run the scorer
    result = await scorer.score(output=sample_text)

    # Display the result
    print("LlamaGuard Scoring Result:")
    print(f"Safe: {result['safe']}")
    print(f"Unsafe Score: {result['extras']['unsafe_score']}")
    if not result['safe']:
        print("Violated Categories:")
        for category, violated in result['extras']['categories'].items():
            if violated:
                print(f" - {category}")


if __name__ == "__main__":
    asyncio.run(main())

import weave
from openai import OpenAI
import pandas as pd
from datasets import load_dataset
import asyncio
import os
import torch
import random
import ast
from weave.scorers.llamaguard_scorer import LlamaGuardScorer
from weave.scorers.moderation_scorer import ToxicityScorer
import os
import json
import gzip
# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"
import weave
from transformers import AutoTokenizer



from datasets import load_dataset
import random

# Initialize scorer
llama_scorer = LlamaGuardScorer(
    model_name_or_path="meta-llama/Llama-Guard-3-8B" if device == "cuda" else "meta-llama/Llama-Guard-3-1B",
    device=device  # Use "cuda" if you have a GPU
)

client = OpenAI(api_key="your api key")
scorer = ToxicityScorer()

weave.init("toxicity_eval")


class PrecisionRecallF1Scorer(weave.Scorer):
    """
    Custom scorer to calculate precision, recall, and F1 at the dataset level.
    """

    @weave.op
    def score(self, label: int, model_output: int) -> dict:
        """
        Compute True Positives, False Positives, and False Negatives for a single row.
        """
        tp = int(label == 1 and model_output == 1)  # True Positive
        fp = int(label == 0 and model_output == 1)  # False Positive
        fn = int(label == 1 and model_output == 0)  # False Negative

        return {"tp": tp, "fp": fp, "fn": fn}

    def summarize(self, score_rows: list) -> dict:
        """
        Summarize precision, recall, and F1 from the row-level scores.
        """
        # Aggregate true positives, false positives, and false negatives
        total_tp = sum(row["tp"] for row in score_rows)
        total_fp = sum(row["fp"] for row in score_rows)
        total_fn = sum(row["fn"] for row in score_rows)

        # Calculate precision, recall, and F1
        precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0
        recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0
        f1 = (
            2 * (precision * recall) / (precision + recall)
            if (precision + recall) > 0
            else 0
        )

        return {
            "precision": precision,
            "recall": recall,
            "f1": f1,
            "tp": total_tp,
            "fp": total_fp,
            "fn": total_fn,
        }



@weave.op
def toxicity_accuracy(label: int, model_output: int) -> dict:
    return {"accuracy": int(model_output == label)}



def load_ds(dataset_name="tcapelle/kaggle-toxic-annotated", split="train", balanced_samples=50, seed=42):
    # Load the dataset from Hugging Face Hub
    print("Loading dataset...")
    dataset = load_dataset(dataset_name, split=split)

    # Convert to Pandas DataFrame for easier manipulation
    df = dataset.to_pandas()

    # Initialize random seed
    random.seed(seed)

    # Define subclasses
    subclasses = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]

    # Remove rows where any subclass column contains -1
    df = df[~df[subclasses].isin([-1]).any(axis=1)]

    # Separate into non-toxic, toxic, and severe toxic categories
    non_toxic_samples = []
    toxic_samples = []
    severe_toxic_samples = []

    for _, row in df.iterrows():
        comment_text = row.get("comment_text", "").strip()
        if not comment_text:
            continue

        # Check if it's non-toxic: all subclass columns must be 0
        if all(row.get(subclass, 0) == 0 for subclass in subclasses):
            non_toxic_samples.append({"text": comment_text, "label": 0})  # Non-toxic
        elif row["severe_toxic"] == 1:
            severe_toxic_samples.append({"text": comment_text, "label": 1})  # Severe toxic
        elif row["toxic"] == 1:
            toxic_samples.append({"text": comment_text, "label": 1})  # Toxic

    # Balance the main categories
    balanced_severe_toxic = random.sample(severe_toxic_samples, min(len(severe_toxic_samples), balanced_samples))
    balanced_toxic = random.sample(toxic_samples, min(len(toxic_samples), balanced_samples))
    balanced_non_toxic = random.sample(non_toxic_samples, min(len(non_toxic_samples), balanced_samples))

    # Combine all samples
    final_data = balanced_severe_toxic + balanced_toxic + balanced_non_toxic
    random.shuffle(final_data)

    return final_data


class LlamaGuard(weave.Model):
    @weave.op
    async def predict(self, text: str) -> int:
        # Initialize the LlamaGuardScorer

        # Run the scorer
        result = await llama_scorer.score(output=text)
        
        # If the text is safe, return 0 (safe)
        if result["safe"]:
            return 0
        return 1        


class OpenAIModerationModel(weave.Model):
    @weave.op
    def predict(self, text: str) -> int:
        response = client.moderations.create(model="omni-moderation-latest", input=text)
        moderation_result = response.results[0]
        return 1 if moderation_result.flagged else 0

class WeaveToxicityScorerModel(weave.Model):
    @weave.op
    def predict(self, text: str) -> int:
        toxicity_result = scorer.score(text)
        return 1 if toxicity_result.get("flagged", False) else 0

# Run evaluations
async def run_evaluations():
    dataset = load_ds()

    print("Dataset loaded...")

    models = {
        "LlamaGuard": LlamaGuard(),
        "OpenAI": OpenAIModerationModel(),
        "toxicityScorer": WeaveToxicityScorerModel(),
    }

    dataset_prepared = [{"text": row["text"], "label": row["label"]} for row in dataset]
    scorers = [toxicity_accuracy, PrecisionRecallF1Scorer()]
    results = {}

    for model_name, model in models.items():
        print(f"\nEvaluating {model_name}...")
        evaluation = weave.Evaluation(
            dataset=dataset_prepared,
            scorers=scorers,
            name=model_name + " Eval"
        )
        results[model_name] = await evaluation.evaluate(model)

if __name__ == "__main__":
    asyncio.run(run_evaluations())