microsoft/hve-core
Publicmirrored from https://github.com/microsoft/hve-coreAvailable
scripts/evals/moderation/moderate.py
197lines · modecode
| 1 | #!/usr/bin/env python3 |
| 2 | # Copyright (c) Microsoft Corporation. |
| 3 | # SPDX-License-Identifier: MIT |
| 4 | """Content moderation CLI using Detoxify toxicity classifier. |
| 5 | |
| 6 | Reads JSON-lines input containing text records, classifies each via Detoxify, |
| 7 | and writes structured JSON output with per-record scores and an overall summary. |
| 8 | Exits with code 1 when any record exceeds the toxicity threshold. |
| 9 | """ |
| 10 | |
| 11 | import argparse |
| 12 | import json |
| 13 | import logging |
| 14 | import sys |
| 15 | from pathlib import Path |
| 16 | from typing import Any, Literal |
| 17 | |
| 18 | EXIT_SUCCESS = 0 |
| 19 | EXIT_FAILURE = 1 |
| 20 | EXIT_ERROR = 2 |
| 21 | |
| 22 | logger = logging.getLogger(__name__) |
| 23 | |
| 24 | |
| 25 | def create_parser() -> argparse.ArgumentParser: |
| 26 | """Create and configure argument parser.""" |
| 27 | parser = argparse.ArgumentParser( |
| 28 | description="Moderate text content using Detoxify toxicity classifier", |
| 29 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 30 | ) |
| 31 | input_group = parser.add_mutually_exclusive_group(required=True) |
| 32 | input_group.add_argument( |
| 33 | "--input", |
| 34 | type=Path, |
| 35 | help="Path to JSON-lines input file with {id, text} records", |
| 36 | ) |
| 37 | input_group.add_argument( |
| 38 | "--stdin", |
| 39 | action="store_true", |
| 40 | help="Read JSON-lines input from stdin", |
| 41 | ) |
| 42 | parser.add_argument( |
| 43 | "--threshold", |
| 44 | type=float, |
| 45 | default=0.5, |
| 46 | help="Toxicity threshold (0.0-1.0); scores above this trigger a flag (default: 0.5)", |
| 47 | ) |
| 48 | parser.add_argument( |
| 49 | "--model", |
| 50 | type=str, |
| 51 | choices=["original", "unbiased", "multilingual"], |
| 52 | default="unbiased", |
| 53 | help="Detoxify model variant (default: unbiased)", |
| 54 | ) |
| 55 | parser.add_argument( |
| 56 | "--output", |
| 57 | type=Path, |
| 58 | required=True, |
| 59 | help="Path to write structured JSON output", |
| 60 | ) |
| 61 | parser.add_argument( |
| 62 | "-v", |
| 63 | "--verbose", |
| 64 | action="store_true", |
| 65 | help="Enable verbose logging", |
| 66 | ) |
| 67 | return parser |
| 68 | |
| 69 | |
| 70 | def configure_logging(verbose: bool = False) -> None: |
| 71 | """Configure logging based on verbosity level.""" |
| 72 | level = logging.DEBUG if verbose else logging.INFO |
| 73 | logging.basicConfig(level=level, format="%(levelname)s: %(message)s") |
| 74 | |
| 75 | |
| 76 | def load_records(input_path: Path | None) -> list[dict[str, str]]: |
| 77 | """Load JSON-lines records from file or stdin.""" |
| 78 | records = [] |
| 79 | source = sys.stdin if input_path is None else input_path.open(encoding="utf-8") |
| 80 | try: |
| 81 | for line_num, line in enumerate(source, start=1): |
| 82 | line = line.strip() |
| 83 | if not line: |
| 84 | continue |
| 85 | try: |
| 86 | record = json.loads(line) |
| 87 | if not isinstance(record, dict): |
| 88 | logger.warning("Line %d: expected object, got %s", line_num, type(record).__name__) |
| 89 | continue |
| 90 | if "id" not in record or "text" not in record: |
| 91 | logger.warning("Line %d: missing required fields (id, text)", line_num) |
| 92 | continue |
| 93 | if not isinstance(record["text"], str): |
| 94 | logger.warning("Line %d: 'text' must be a string, got %s", line_num, type(record["text"]).__name__) |
| 95 | continue |
| 96 | records.append(record) |
| 97 | except json.JSONDecodeError as e: |
| 98 | logger.warning("Line %d: JSON parse error: %s", line_num, e) |
| 99 | finally: |
| 100 | if input_path is not None: |
| 101 | source.close() |
| 102 | logger.info("Loaded %d records", len(records)) |
| 103 | return records |
| 104 | |
| 105 | |
| 106 | def classify_records( |
| 107 | records: list[dict[str, str]], |
| 108 | model_name: Literal["original", "unbiased", "multilingual"], |
| 109 | threshold: float, |
| 110 | ) -> list[dict[str, Any]]: |
| 111 | """Classify records using Detoxify and return results with flag status.""" |
| 112 | try: |
| 113 | from detoxify import Detoxify |
| 114 | except ImportError: |
| 115 | logger.error("detoxify package not installed; run: uv pip install -r requirements.txt") |
| 116 | sys.exit(EXIT_ERROR) |
| 117 | |
| 118 | logger.info("Loading Detoxify model: %s", model_name) |
| 119 | model = Detoxify(model_name) |
| 120 | |
| 121 | results = [] |
| 122 | for record in records: |
| 123 | record_id = record["id"] |
| 124 | text = record["text"] |
| 125 | logger.debug("Classifying record: %s", record_id) |
| 126 | |
| 127 | scores = model.predict(text) |
| 128 | # Convert numpy types to native Python floats |
| 129 | scores = {k: float(v) for k, v in scores.items()} |
| 130 | |
| 131 | flagged_labels = [label for label, score in scores.items() if score > threshold] |
| 132 | flagged = len(flagged_labels) > 0 |
| 133 | |
| 134 | results.append( |
| 135 | { |
| 136 | "id": record_id, |
| 137 | "scores": scores, |
| 138 | "flagged": flagged, |
| 139 | "flaggedLabels": flagged_labels, |
| 140 | } |
| 141 | ) |
| 142 | if flagged: |
| 143 | logger.warning( |
| 144 | "Record %s FLAGGED: %s", |
| 145 | record_id, |
| 146 | ", ".join(f"{label}={scores[label]:.3f}" for label in flagged_labels), |
| 147 | ) |
| 148 | |
| 149 | return results |
| 150 | |
| 151 | |
| 152 | def write_output(results: list[dict[str, Any]], output_path: Path) -> None: |
| 153 | """Write structured JSON output with per-record results and summary.""" |
| 154 | flagged_count = sum(1 for r in results if r["flagged"]) |
| 155 | output = { |
| 156 | "records": results, |
| 157 | "summary": { |
| 158 | "total": len(results), |
| 159 | "flaggedCount": flagged_count, |
| 160 | }, |
| 161 | } |
| 162 | output_path.parent.mkdir(parents=True, exist_ok=True) |
| 163 | output_path.write_text(json.dumps(output, indent=2), encoding="utf-8") |
| 164 | logger.info("Wrote output to %s", output_path) |
| 165 | |
| 166 | |
| 167 | def main() -> int: |
| 168 | """Main entry point.""" |
| 169 | parser = create_parser() |
| 170 | args = parser.parse_args() |
| 171 | configure_logging(args.verbose) |
| 172 | |
| 173 | if args.threshold < 0.0 or args.threshold > 1.0: |
| 174 | logger.error("Threshold must be between 0.0 and 1.0") |
| 175 | return EXIT_ERROR |
| 176 | |
| 177 | input_path = args.input |
| 178 | records = load_records(input_path) |
| 179 | if not records: |
| 180 | logger.warning("No records to process") |
| 181 | write_output([], args.output) |
| 182 | return EXIT_SUCCESS |
| 183 | |
| 184 | results = classify_records(records, args.model, args.threshold) |
| 185 | write_output(results, args.output) |
| 186 | |
| 187 | flagged_count = sum(1 for r in results if r["flagged"]) |
| 188 | if flagged_count > 0: |
| 189 | logger.error("Content moderation failed: %d/%d records flagged", flagged_count, len(results)) |
| 190 | return EXIT_FAILURE |
| 191 | |
| 192 | logger.info("Content moderation passed: all %d records clean", len(results)) |
| 193 | return EXIT_SUCCESS |
| 194 | |
| 195 | |
| 196 | if __name__ == "__main__": |
| 197 | sys.exit(main()) |