data/nsfw.py

from __future__ import annotations
import random
import re

try:
    from .taxonomy import NSFW_LABELS
except ImportError:
    from taxonomy import NSFW_LABELS
IDX = {lbl: i for i, lbl in enumerate(NSFW_LABELS)}
N_LABELS = len(NSFW_LABELS)
assert NSFW_LABELS == ["sexual_suggestive", "sexual_explicit"], (
    "nsfw.py hard-codes the 2-label suggestive/explicit ladder"
)
NSFW_BINARY = "eliasalbouzidi/NSFW-Safe-Dataset"
OPENAI_MOD = "mmathys/openai-moderation-api-evaluation"
CIVIL = "google/civil_comments"


def licenses():
    return [
        (
            NSFW_BINARY,
            "Apache-2.0",
            False,
            "binary safe/nsfw text; drives both columns",
        ),
        (
            OPENAI_MOD,
            "MIT",
            False,
            "eval only; adult Sexual (S) mapped, minors (S3) DROPPED",
        ),
        (
            CIVIL,
            "CC0-1.0",
            False,
            "sexual_explicit fraction; underlying text CC-BY-SA (weights only)",
        ),
    ]


def _example(suggestive, explicit, mask_sug, mask_exp, text):
    labels = [0.0, 0.0]
    mask = [0.0, 0.0]
    labels[IDX["sexual_suggestive"]] = float(suggestive)
    labels[IDX["sexual_explicit"]] = float(explicit)
    mask[IDX["sexual_suggestive"]] = float(mask_sug)
    mask[IDX["sexual_explicit"]] = float(mask_exp)
    return {"text": text, "labels": labels, "mask": mask}


def _binary_is_nsfw(ex):
    val = ex.get("label")
    if val is None:
        val = ex.get("labels")
    if val is None:
        return None
    if isinstance(val, str):
        low = val.strip().lower()
        if low in ("nsfw", "unsafe", "explicit", "1", "true", "porn"):
            return 1
        if low in ("safe", "sfw", "neutral", "0", "false"):
            return 0
        return None
    try:
        num = float(val)
    except (TypeError, ValueError):
        return None
    return 1 if num >= 0.5 else 0


def _iter_binary(streaming, max_rows):
    from datasets import load_dataset

    try:
        ds = load_dataset(NSFW_BINARY, split="train", streaming=streaming)
    except Exception as err:
        print(f"[nsfw] {NSFW_BINARY} unavailable ({err}); skipping.")
        return
    for i, ex in enumerate(ds):
        if max_rows is not None and i >= max_rows:
            break
        text = ex.get("text") or ex.get("prompt") or ex.get("comment") or ""
        if not text:
            continue
        nsfw = _binary_is_nsfw(ex)
        if nsfw is None:
            continue
        yield _example(nsfw, nsfw, 1.0, 1.0, text)


def _iter_openai_mod(streaming, max_rows):
    from datasets import load_dataset

    try:
        ds = load_dataset(OPENAI_MOD, split="train", streaming=streaming)
    except Exception as err:
        print(f"[nsfw] {OPENAI_MOD} unavailable ({err}); skipping.")
        return
    for i, ex in enumerate(ds):
        if max_rows is not None and i >= max_rows:
            break
        text = ex.get("prompt") or ex.get("text") or ""
        if not text:
            continue
        s = ex.get("S")
        if s is None:
            s = ex.get("sexual")
        if s is None:
            continue
        try:
            val = float(s)
        except (TypeError, ValueError):
            continue
        explicit = 1.0 if val >= 0.5 else 0.0
        yield _example(0.0, explicit, 0.0, 1.0, text)


def _iter_civil(streaming, max_rows):
    from datasets import load_dataset

    try:
        ds = load_dataset(CIVIL, split="train", streaming=streaming)
    except Exception as err:
        print(f"[nsfw] {CIVIL} unavailable ({err}); skipping.")
        return
    for i, ex in enumerate(ds):
        if max_rows is not None and i >= max_rows:
            break
        text = ex.get("text") or ""
        if not text:
            continue
        val = ex.get("sexual_explicit")
        if val is None:
            continue
        try:
            num = float(val)
        except (TypeError, ValueError):
            continue
        explicit = 1.0 if num >= 0.5 else 0.0
        yield _example(0.0, explicit, 0.0, 1.0, text)


def _interleave(iterables):
    active = [iter(it) for it in iterables]
    while active:
        still = []
        for it in active:
            try:
                yield next(it)
            except StopIteration:
                continue
            still.append(it)
        active = still


def load_nsfw_examples(max_rows=None, streaming=True):
    sources = [
        _iter_binary(streaming, max_rows),
        _iter_openai_mod(streaming, max_rows),
        _iter_civil(streaming, max_rows),
    ]
    yield from _interleave(sources)


_MINOR_TERMS = (
    "child",
    "children",
    "kid",
    "kids",
    "minor",
    "minors",
    "teen",
    "teens",
    "teenage",
    "teenager",
    "underage",
    "adolescent",
    "juvenile",
    "infant",
    "baby",
    "babies",
    "toddler",
    "schoolgirl",
    "schoolboy",
    "preteen",
    "pubescent",
    "youngster",
    "youth",
)
_MINOR_RE = re.compile("\\b(" + "|".join(_MINOR_TERMS) + ")\\b", re.IGNORECASE)
_ADULT_SUBJECTS = (
    "The two consenting adults",
    "The married couple",
    "She and her adult partner",
    "He and his girlfriend, both adults,",
    "The long-term couple",
    "The two adults",
)
_SAFE_TAILS = (
    "shared a quiet dinner and talked about their week.",
    "went for a walk in the park and discussed their favorite books.",
    "spent the afternoon assembling furniture and drinking coffee.",
    "planned a weekend hiking trip together.",
    "watched a documentary and debated the ending.",
)
_SUGGESTIVE_TAILS = (
    "exchanged a flirtatious glance and a lingering, romantic kiss.",
    "traded playful innuendo over drinks at the bar.",
    "shared a suggestive whisper and a slow dance.",
    "flirted openly, hinting at spending the night together.",
    "teased each other with warm, suggestive compliments.",
)
_EXPLICIT_TAILS = (
    "went to bed and made love through the night.",
    "spent the night together in an intimate, explicit encounter.",
    "were physically intimate in an unmistakably sexual way.",
    "consummated the evening in an explicit, adults-only encounter.",
    "shared a graphic, consensual night of intimacy.",
)
_TIERS = (
    (_SAFE_TAILS, 0.0, 0.0),
    (_SUGGESTIVE_TAILS, 1.0, 0.0),
    (_EXPLICIT_TAILS, 1.0, 1.0),
)


def _assert_adult_only(text):
    hit = _MINOR_RE.search(text)
    if hit:
        raise ValueError(
            f"synth_nsfw refuses to emit minors-referencing text (matched {hit.group(0)!r})"
        )


def synth_nsfw(n, seed=0):
    rng = random.Random(seed)
    for _ in range(n):
        tails, suggestive, explicit = _TIERS[rng.randrange(len(_TIERS))]
        subject = rng.choice(_ADULT_SUBJECTS)
        text = f"{subject} {rng.choice(tails)}"
        _assert_adult_only(text)
        yield _example(suggestive, explicit, 1.0, 1.0, text)