task/task_05/step4_bias_audit.py

"""
step4_bias_audit.py
====================
Task 5 — Component 4: Stereotype / bias detection in generated captions.

Implements a lexicon-based bias audit across three demographic axes:
  - Gender stereotypes (subject + attribute co-occurrence)
  - Age stereotypes
  - Occupation bias

The method:
  For each caption, we check whether it contains a SUBJECT term from a
  demographic group AND an ATTRIBUTE term that is stereotypically associated
  with that group. A caption is flagged if a (subject, attribute) co-occurrence
  is found in the same caption.

This produces a demographic_group → attribute_frequency table — the
equivalent of a "bias heatmap".

Public API
----------
    audit_caption(caption) -> dict
    run_bias_audit(records, save_dir)  -> list[dict]
    _load_or_use_precomputed(save_dir) -> list[dict]

Standalone usage
----------------
    export PYTHONPATH=.
    venv/bin/python task/task_05/step4_bias_audit.py
"""

import os
import sys
import json
import re
from collections import defaultdict

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))


# ─────────────────────────────────────────────────────────────────────────────
# Stereotype lexicon
# ─────────────────────────────────────────────────────────────────────────────

# Structure: {group_name: {subjects: [...], attributes: [...]}}
STEREOTYPE_LEXICON = {
    "gender_women_domestic": {
        "subjects":    ["woman", "women", "she", "her", "female", "lady", "girl", "mother"],
        "attributes":  ["cooking", "cleaning", "housework", "baking", "shopping",
                        "knitting", "laundry", "ironing", "sewing", "nanny",
                        "caring for", "taking care"],
        "label": "Women → Domestic roles",
    },
    "gender_men_sports": {
        "subjects":    ["man", "men", "he", "his", "male", "guy", "boy", "father"],
        "attributes":  ["sports", "football", "basketball", "soccer", "driving",
                        "fighting", "competing", "wrestling", "rugby", "racing",
                        "construction", "grilling", "barbecue"],
        "label": "Men → Sports / Physical roles",
    },
    "gender_women_nursing": {
        "subjects":    ["woman", "women", "female", "nurse"],
        "attributes":  ["nurse", "nursing", "caring", "attendant", "aid",
                        "receptionist"],
        "label": "Women → Healthcare support roles",
    },
    "gender_men_leadership": {
        "subjects":    ["man", "men", "male", "doctor", "officer"],
        "attributes":  ["doctor", "boss", "officer", "engineer", "pilot",
                        "executive", "manager", "director"],
        "label": "Men → Leadership / Technical roles",
    },
    "age_elderly_negative": {
        "subjects":    ["elderly", "old", "older", "senior", "aged", "elder"],
        "attributes":  ["frail", "weak", "slow", "confused", "alone", "sitting",
                        "resting", "feeble"],
        "label": "Elderly → Negative / passive attributes",
    },
    "age_young_reckless": {
        "subjects":    ["young", "youth", "teenager", "kid", "teen"],
        "attributes":  ["reckless", "running", "jumping", "skateboarding",
                        "rowdy", "skateboard", "bike"],
        "label": "Young → Reckless / energetic",
    },
}


# ─────────────────────────────────────────────────────────────────────────────
# Audit functions
# ─────────────────────────────────────────────────────────────────────────────

def _tokenize(text: str) -> set:
    """Lowercase word-tokenize."""
    return set(re.findall(r"\b[a-z]+\b", text.lower()))


def audit_caption(caption: str) -> dict:
    """
    Check a single caption for stereotype patterns.

    Returns:
        {
          flagged: bool,
          matches: list of {group, label, subject, attribute}
        }
    """
    words   = _tokenize(caption)
    matches = []

    for group, info in STEREOTYPE_LEXICON.items():
        subjects_hit    = words & set(info["subjects"])
        attributes_hit  = words & set(info["attributes"])

        if subjects_hit and attributes_hit:
            for subj in subjects_hit:
                for attr in attributes_hit:
                    matches.append({
                        "group":     group,
                        "label":     info["label"],
                        "subject":   subj,
                        "attribute": attr,
                    })

    return {
        "flagged": len(matches) > 0,
        "matches": matches,
    }


def run_bias_audit(records: list,
                   save_dir: str = "task/task_05/results") -> list:
    """
    Audit all captions for stereotypes.

    Args:
        records : list of {image_id, caption, ...} (from step2 or step3)
        save_dir: where to save bias_audit.json

    Returns:
        list of {image_id, caption, flagged, matches, group_flags}
    """
    print("=" * 68)
    print("  Task 5 — Step 4: Bias / Stereotype Audit")
    print(f"  Auditing {len(records)} captions ...")
    print("=" * 68)

    results = []
    group_counts = defaultdict(int)

    for rec in records:
        audit = audit_caption(rec["caption"])
        for m in audit["matches"]:
            group_counts[m["group"]] += 1

        results.append({
            "image_id":   rec.get("image_id", 0),
            "caption":    rec["caption"],
            "flagged":    audit["flagged"],
            "matches":    audit["matches"],
            "group_flags": list({m["group"] for m in audit["matches"]}),
        })

    # Frequency table
    freq_table = {}
    for group, cnt in group_counts.items():
        label  = STEREOTYPE_LEXICON[group]["label"]
        freq_table[group] = {
            "label": label,
            "count": cnt,
            "rate":  round(cnt / max(len(records), 1), 4),
        }

    os.makedirs(save_dir, exist_ok=True)
    audit_path = os.path.join(save_dir, "bias_audit.json")
    with open(audit_path, "w") as f:
        json.dump({
            "records":    results,
            "freq_table": freq_table,
        }, f, indent=2)

    n_flagged = sum(1 for r in results if r["flagged"])
    print(f"  Flagged captions  : {n_flagged} ({100*n_flagged/max(len(results),1):.1f}%)")
    print(f"  Stereotype frequencies:")
    for g, info in sorted(freq_table.items(), key=lambda x: -x[1]["count"]):
        print(f"    {info['label']:45s}  count={info['count']:4d}  rate={info['rate']:.3f}")
    print(f"  Saved -> {audit_path}")
    return results, freq_table


# ─────────────────────────────────────────────────────────────────────────────
# Load / create precomputed
# ─────────────────────────────────────────────────────────────────────────────

def _load_or_use_precomputed(save_dir: str) -> tuple:
    """Return cached bias_audit.json or compute from cached captions."""
    cache = os.path.join(save_dir, "bias_audit.json")
    if os.path.exists(cache):
        with open(cache) as f:
            data = json.load(f)
        print(f"  OK  Loaded cached bias audit from {cache}")
        return data["records"], data["freq_table"]

    # Compute from precomputed captions
    from step2_prepare_data import _load_or_use_precomputed as load_caps
    caps = load_caps(save_dir)
    return run_bias_audit(caps, save_dir)


# ─────────────────────────────────────────────────────────────────────────────
# Standalone
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    SAVE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")
    records, freq_table = _load_or_use_precomputed(SAVE_DIR)
    flagged = [r for r in records if r["flagged"]]
    print(f"\nSample flagged captions:")
    for r in flagged[:5]:
        print(f"  \"{r['caption']}\"")
        for m in r["matches"][:2]:
            print(f"    -> {m['label']}  [{m['subject']} + {m['attribute']}]")