ViToxReduce Pipeline

A comprehensive pipeline for reducing toxicity in text using a streamlined 3-tier architecture.

Pipeline Architecture

📋 Description

ViToxReduce is an automated system for reducing toxicity in text using a 3-tier architecture:

Safety Checker: Detects toxic sentences
Span Tagger: Locates toxic phrases
Contextual Rewriter: Safely rewrites sentences

🏗️ System Architecture

The ViToxReduce pipeline uses a streamlined 3-tier architecture to efficiently process and detoxify text:

Pipeline Architecture

Example Flow

The following diagram shows an example of how the pipeline processes a toxic sentence:

Example Flow

Tier 1: Safety Checker

Function: Classifies input sentences as safe (safe) or unsafe (toxic)
Model: PhoBERT (Sequence Classification)
Processing Flow:
- If Safe → Returns the original sentence unchanged, terminates processing (saves computational resources for ~50% of non-toxic sentences)
- If Unsafe → Proceeds to Tier 2 for further processing

Tier 2: Span Tagger

Function: Scans sentences and labels each word to identify toxic phrase locations (Toxic Spans)
Model: PhoBERT (Token Classification / BIO Tagging)
Output: A list of spans S = [s_1, s_2, ...] (character indices)
Role: Provides "hints" to Tier 3 about which parts to focus on, preventing over-editing of safe regions

Tier 3: Contextual Rewriter

Function: Generates new sentences that remove or replace toxic spans while preserving context and intent of remaining parts
Model: BARTpho (Seq2Seq - Span-guided)
Input: Original sentence X + Span list S (combined into prompt: "Original sentence: X. Words to fix: S")
Generation Mechanism: Uses One-time Rewriting with Beam Search to select the best result (Top-1) without iterative checking, improving processing speed

📦 Installation

System Requirements

Python >= 3.7
CUDA (optional, for GPU acceleration)

Install via pip

pip install vitoxreduce

(Dev mode) From source:

pip install -r requirements.txt
pip install -e .

Required Models

You need 3 fine-tuned models (can be loaded directly from Hugging Face Hub):

BARTpho Rewriter (Tier 3) — joshswift/bartpho-rewriter
PhoBERT Span Locator (Tier 2) — joshswift/phobert-span
PhoBERT Toxicity Classifier (Tier 1) — joshswift/phobert-toxicity

You can use the repo IDs directly (online), or download locally:

pip install huggingface_hub
huggingface-cli login          # if private or rate-limited
huggingface-cli download joshswift/bartpho-rewriter   --local-dir ./models/rewriter
huggingface-cli download joshswift/phobert-span       --local-dir ./models/span
huggingface-cli download joshswift/phobert-toxicity   --local-dir ./models/toxicity

Dataset (optional, for testing/eval)

Hugging Face dataset: joshswift/vitoxrewrite
Download locally (JSONL kept intact):

pip install huggingface_hub
huggingface-cli login  # if needed
huggingface-cli download joshswift/vitoxrewrite --local-dir ./dataset

After download, you will have:

dataset/vitoxrewrite_train.jsonl
dataset/vitoxrewrite_validation.jsonl
dataset/vitoxrewrite_test.jsonl

🚀 Quick Start

Option 1: Automated Setup Script (Recommended for First-Time Users)

The easiest way to get started is using the automated setup script that installs the package, downloads models, and runs a test:

# Clone or download the repository
cd vitoxreduce_pipeline_github

# Run the automated setup script
python examples/example_usage.py

What the script does:

Installs vitoxreduce package from PyPI
Downloads 3 required models from Hugging Face (skips if already exists)
Runs a smoke test with sample text
Saves results to ./results/smoke_test_TIMESTAMP.json

Customize the script:

# Use custom model directory
python examples/example_usage.py --models-dir ./my_models

# Use Hugging Face token (if repos are private/rate-limited)
python examples/example_usage.py --token hf_xxxxxxxxxxxxx

# Custom output file
python examples/example_usage.py --output ./my_results.json

Option 2: Manual Setup (Step-by-Step)

If you prefer to set up manually or need more control, follow these steps:

1) Install Package

pip install vitoxreduce

2) Download Models (Optional - can use online models instead)

pip install huggingface_hub
huggingface-cli download joshswift/bartpho-rewriter   --local-dir ./models/rewriter
huggingface-cli download joshswift/phobert-span       --local-dir ./models/span
huggingface-cli download joshswift/phobert-toxicity   --local-dir ./models/toxicity

3) Run CLI

# Single sentence (using online models)
vitoxreduce \
  --input "Từ lúc mấy bro cmt cực kì cl gì đấy..." \
  --rewriter_model joshswift/bartpho-rewriter \
  --span_locator_model joshswift/phobert-span \
  --toxicity_detector_model joshswift/phobert-toxicity \
  --output result.json \
  --verbose

# Or use local models
vitoxreduce \
  --input "Your text here" \
  --rewriter_model ./models/rewriter \
  --span_locator_model ./models/span \
  --toxicity_detector_model ./models/toxicity \
  --output result.json \
  --verbose

4) Python API

from vitoxreduce import ViToxReducePipeline

pipeline = ViToxReducePipeline(
    rewriter_model_path="joshswift/bartpho-rewriter",  # or local path
    span_locator_model_path="joshswift/phobert-span",
    toxicity_detector_model_path="joshswift/phobert-toxicity",
    num_beams=5,
)

result = pipeline.process("Your text here", verbose=True)
print("Rewritten:", result["rewritten"])
print("Safe?:", result["rewritten_is_safe"])
print("Toxicity:", result["toxicity_score"], "→", result["rewritten_toxicity_score"])

⚙️ Command-Line Arguments

Required Arguments

--rewriter_model: Path to BARTpho rewriter model (Tier 3) or Hugging Face repo ID
--span_locator_model: Path to PhoBERT span locator model (Tier 2) or Hugging Face repo ID
--toxicity_detector_model: Path to PhoBERT toxicity classifier (Tier 1) or Hugging Face repo ID
--input: Input text, text file (one sentence per line), or JSONL file with 'comment' field

Optional Arguments

--output: Output JSON file path (default: auto-generated or stdout)
--num_beams: Number of beams for beam search (default: 5)
--toxicity_threshold: Threshold for unsafe classification (default: 0.5)
--device: Device to use - cuda, cpu, or auto (default: auto)
--verbose: Print detailed processing information
--mode: Processing mode - auto, single, file, or jsonl (default: auto)

📊 Output Format

The pipeline returns a dictionary with the following keys:

{
    "original": "Original sentence",
    "is_safe": False,                    # True if safe, False if unsafe
    "toxicity_score": 0.85,              # Toxicity score (0.0-1.0)
    "spans": [(10, 15), (20, 25)],       # List of detected spans (character indices)
    "span_texts": ["toxic", "word"],     # List of span texts
    "rewritten": "Rewritten sentence",   # Rewritten sentence (or original if safe)
    "rewritten_is_safe": True,           # True if rewritten sentence is safe
    "rewritten_toxicity_score": 0.15,    # Toxicity score of rewritten sentence
    "processing_time": 1.2                # Processing time (seconds)
}

📁 Project Structure

vitoxreduce_pipeline_github/
├── vitoxreduce/                    # Main package
│   ├── pipeline.py                 # Main 3-tier pipeline
│   ├── tier1_toxicity_detector.py # Tier 1: Safety Checker
│   ├── tier2_span_locator.py      # Tier 2: Span Tagger
│   ├── tier3_rewrite_generator.py # Tier 3: Contextual Rewriter
│   └── ...
├── scripts/                        # CLI scripts
│   └── run_pipeline.py            # Main CLI entry point
├── examples/                       # Usage examples
│   └── example_usage.py           # Automated setup & test script
├── dataset/                        # Sample dataset (optional)
└── requirements.txt                # Python dependencies

🔍 Evaluation Metrics

The pipeline supports the following metrics (when reference is available):

BLEU: Measures similarity with reference
SIM: Semantic similarity (using SBERT)
FL: Fluency score (based on GPT-2 PPL)
STA: Toxicity Drop (toxicity reduction)
J-score: Joint score = SIM × FL × normalized_tox_drop

Performance Comparison

The above chart shows ViToxReduce's performance compared to baseline methods across key metrics including BLEU, SIM, J-score, and toxicity reduction.

📖 API Documentation