social-data-attribution-w2/scripts/diagnostics/probe_fasttext_quality.py

"""Phase 1: Empirical FastText quality model verification.

Part A: Determine which label (__label__0 vs __label__1) the model
assigns to unambiguously high vs low quality text.

Part B: Test whether truncated text receives higher quality scores
than full-length text from the same source passage.

No R2 credentials required. Runs locally.
"""

from __future__ import annotations

import json
import logging
from dataclasses import asdict, dataclass, field
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger(__name__)

LABEL_VERIFICATION_CASES: list[dict[str, str]] = [
    {
        "name": "wikipedia_paragraph",
        "expected": "high",
        "text": (
            "The mitochondrion is a double-membrane-bound organelle found in most "
            "eukaryotic organisms. Mitochondria generate most of the cell's supply "
            "of adenosine triphosphate (ATP), used as a source of chemical energy. "
            "They were first discovered by Albert von Kolliker in 1857 and the term "
            "mitochondrion was coined by Carl Benda in 1898. The organelle is composed "
            "of compartments that carry out specialized functions, including the outer "
            "membrane, intermembrane space, inner membrane, cristae, and matrix."
        ),
    },
    {
        "name": "news_article",
        "expected": "high",
        "text": (
            "Federal Reserve officials voted unanimously to hold the benchmark interest "
            "rate steady at a range of 5.25% to 5.50%, citing continued progress on "
            "inflation alongside a resilient labor market. Chair Jerome Powell noted "
            "that policymakers need greater confidence that inflation is moving "
            "sustainably toward the 2% target before considering rate cuts. The "
            "decision was widely anticipated by financial markets, with futures "
            "pricing reflecting expectations of the first cut later in the year."
        ),
    },
    {
        "name": "academic_abstract",
        "expected": "high",
        "text": (
            "We present a novel framework for causal inference in observational studies "
            "with high-dimensional confounders. Our method combines doubly robust "
            "estimation with targeted regularization to achieve root-n consistent "
            "estimation of average treatment effects. Through extensive simulations "
            "and an application to electronic health records, we demonstrate that our "
            "approach maintains nominal coverage while substantially reducing mean "
            "squared error compared to existing methods. The theoretical guarantees "
            "hold under model misspecification of either the outcome or propensity model."
        ),
    },
    {
        "name": "spam_viagra",
        "expected": "low",
        "text": (
            "BUY NOW CHEAP VIAGRA CIALIS ONLINE!!! Best prices guaranteed click here "
            "www.bestpills.xyz FREE SHIPPING on all orders!!! Act now limited time "
            "offer discount 90% off!!!! Order today get bonus pills FREE!!!!"
        ),
    },
    {
        "name": "gibberish",
        "expected": "low",
        "text": (
            "asdf jkl; qwerty zxcvbn 12345 67890 aaa bbb ccc ddd eee fff ggg "
            "hhh iii jjj kkk lll mmm nnn ooo ppp qqq rrr sss ttt uuu vvv www "
            "xxx yyy zzz 0000 1111 2222 3333 4444 5555 6666 7777 8888 9999"
        ),
    },
    {
        "name": "seo_spam",
        "expected": "low",
        "text": (
            "Best cheap hotels in New York City 2024. Looking for best cheap hotels "
            "in New York City? Our guide to best cheap hotels in New York City has "
            "the best cheap hotels in New York City for your next trip. Book the "
            "best cheap hotels in New York City today! Best cheap hotels New York "
            "City reviews. Best cheap hotels New York City deals."
        ),
    },
]

TRUNCATION_PASSAGES: list[dict[str, str]] = [
    {
        "name": "encyclopedic",
        "text": (
            "The Great Barrier Reef is the world's largest coral reef system, "
            "composed of over 2,900 individual reefs and 900 islands stretching "
            "for over 2,300 kilometres over an area of approximately 344,400 "
            "square kilometres. The reef is located in the Coral Sea, off the "
            "coast of Queensland, Australia. It can be seen from outer space and "
            "is the world's biggest single structure made by living organisms. "
            "This reef structure is composed of and built by billions of tiny "
            "organisms, known as coral polyps. The Great Barrier Reef supports "
            "a wide diversity of life and was selected as a World Heritage Site "
            "in 1981. CNN labelled it one of the seven natural wonders of the "
            "world. The Queensland National Trust named it a state icon of "
            "Queensland. A large part of the reef is protected by the Great "
            "Barrier Reef Marine Park, which helps to limit the impact of human "
            "use, such as fishing and tourism. Other environmental pressures "
            "on the reef and its ecosystem include runoff, climate change "
            "accompanied by mass coral bleaching, and cyclic population "
            "outbreaks of the crown-of-thorns starfish. According to a study "
            "published in October 2012 by the Proceedings of the National "
            "Academy of Sciences, the reef has lost more than half its coral "
            "cover since 1985."
        ),
    },
    {
        "name": "news_report",
        "text": (
            "Authorities in coastal regions have issued evacuation orders as "
            "Hurricane Maria strengthens to Category 4 status with maximum "
            "sustained winds of 155 mph. The National Hurricane Center projects "
            "the storm will make landfall within 48 hours, potentially bringing "
            "catastrophic storm surge of 12 to 18 feet along vulnerable "
            "shoreline communities. Emergency management officials have opened "
            "135 shelters across three counties and mobilized National Guard "
            "units to assist with evacuations. Power companies are pre-staging "
            "repair crews and equipment in anticipation of widespread outages "
            "that could last weeks in heavily impacted areas. Governor Mitchell "
            "declared a state of emergency and urged all residents in evacuation "
            "zones to leave immediately. Schools, government offices, and "
            "non-essential businesses in the projected path have been ordered "
            "closed through the end of the week. Fuel shortages have been "
            "reported at gas stations along major evacuation routes."
        ),
    },
    {
        "name": "forum_post",
        "text": (
            "Hey everyone, I just got my new gaming PC set up and I'm having "
            "issues with the GPU drivers. Every time I try to update to the "
            "latest version, the screen goes black and I have to restart. Has "
            "anyone else experienced this? I'm running Windows 11 with an RTX "
            "4070 and the latest BIOS update. I tried DDU in safe mode but the "
            "problem persists. My temps are fine, around 65C under load. PSU "
            "is a Corsair 850W so power shouldn't be an issue. Any suggestions "
            "would be appreciated. I've been dealing with this for a week now "
            "and it's driving me crazy. Already submitted a support ticket to "
            "NVIDIA but haven't heard back yet. Also tried an older driver "
            "version and same thing happens."
        ),
    },
    {
        "name": "product_page",
        "text": (
            "Premium Stainless Steel Water Bottle - 32oz Double Wall Vacuum "
            "Insulated. Keeps drinks cold for 24 hours and hot for 12 hours. "
            "BPA-free, leak-proof lid with built-in carrying loop. Available "
            "in 12 colors. Features a powder-coated exterior for extra grip "
            "and durability. Wide mouth opening fits standard ice cubes and "
            "is easy to clean. Compatible with most car cup holders. Perfect "
            "for gym, office, hiking, and everyday use. Each bottle comes with "
            "a lifetime warranty against manufacturing defects. Our bottles are "
            "made from food-grade 18/8 stainless steel with no plastic liner "
            "contact with your beverage. Free shipping on orders over $25. "
            "Rated 4.8 stars from over 15,000 customer reviews. See our FAQ "
            "for care instructions and replacement parts."
        ),
    },
    {
        "name": "spam_with_content",
        "text": (
            "AMAZING DEAL! Don't miss out on this incredible opportunity to "
            "earn $5000 per week working from home! No experience needed! "
            "Our proven system has helped thousands of people achieve financial "
            "freedom. Sign up today and get instant access to our exclusive "
            "training materials. Limited spots available! This offer expires "
            "soon so act now! Click the link below to get started on your "
            "journey to wealth. Remember, the sooner you start, the sooner "
            "you'll see results. Our top earners are making over $20,000 per "
            "month. Join the revolution today! Special bonus for new members "
            "who sign up in the next 24 hours."
        ),
    },
]

BOILERPLATE_SUFFIXES = [
    "\n\nRead more...",
    "\n\nCookie Policy: This website uses cookies to improve your experience. "
    "By continuing to browse the site, you agree to our use of cookies. "
    "Accept | Decline | Learn More",
    "\n\n\xa9 2024 All rights reserved. Terms of Service | Privacy Policy | "
    "Contact Us | About | Sitemap | Accessibility",
]


@dataclass
class ProbeResult:
    part_a: list[dict[str, object]] = field(default_factory=list)
    part_b_truncation: list[dict[str, object]] = field(default_factory=list)
    part_b_length_curve: list[dict[str, object]] = field(default_factory=list)
    label_map: dict[str, object] = field(default_factory=dict)
    raw_model_labels: list[str] = field(default_factory=list)


def run_part_a(model, label_map) -> list[dict[str, object]]:
    log.info("\n=== Part A: Label direction verification ===\n")
    results: list[dict[str, object]] = []

    for case in LABEL_VERIFICATION_CASES:
        clean = case["text"].replace("\n", " ").strip()
        labels_raw, probs_raw = model.predict(clean, k=-1)
        raw_output = list(zip(labels_raw, [float(p) for p in probs_raw]))

        result: dict[str, object] = {
            "name": case["name"],
            "expected_quality": case["expected"],
            "raw_model_output": raw_output,
            "label_1_prob": None,
            "label_0_prob": None,
        }
        for label, prob in raw_output:
            if label == "__label__1":
                result["label_1_prob"] = float(prob)
            elif label == "__label__0":
                result["label_0_prob"] = float(prob)

        log.info(
            "%-25s expected=%-4s  __label__1=%.4f  __label__0=%.4f",
            case["name"],
            case["expected"],
            result["label_1_prob"] or 0,
            result["label_0_prob"] or 0,
        )
        results.append(result)

    high_cases = [r for r in results if r["expected_quality"] == "high"]
    low_cases = [r for r in results if r["expected_quality"] == "low"]
    avg_label1_high = sum(r["label_1_prob"] for r in high_cases) / len(high_cases)
    avg_label1_low = sum(r["label_1_prob"] for r in low_cases) / len(low_cases)

    log.info("\nAvg P(__label__1) for HIGH quality cases: %.4f", avg_label1_high)
    log.info("Avg P(__label__1) for LOW quality cases:  %.4f", avg_label1_low)

    if avg_label1_high > avg_label1_low:
        log.info("RESULT: __label__1 = high quality (our assumption is CORRECT)")
    else:
        log.info("RESULT: __label__1 = low quality (LABELS ARE INVERTED)")

    return results


def score_text(model, text: str) -> float:
    clean = text.replace("\n", " ").strip()
    labels_raw, probs_raw = model.predict(clean, k=-1)
    for label, prob in zip(labels_raw, probs_raw):
        if label == "__label__1":
            return float(prob)
    return 0.0


def truncate_words(text: str, word_count: int) -> str:
    words = text.split()
    if len(words) <= word_count:
        return text
    return " ".join(words[:word_count])


def run_part_b(model) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
    log.info("\n=== Part B: Truncation behavior probe ===\n")
    truncation_results: list[dict[str, object]] = []
    length_curve_results: list[dict[str, object]] = []

    for passage in TRUNCATION_PASSAGES:
        words = passage["text"].split()
        total_words = len(words)
        full_score = score_text(model, passage["text"])

        mid_cutoff = total_words // 2
        truncated_mid = truncate_words(passage["text"], mid_cutoff)
        truncated_mid_score = score_text(model, truncated_mid)

        boilerplate_text = passage["text"] + BOILERPLATE_SUFFIXES[0]
        boilerplate_truncated = truncate_words(passage["text"], mid_cutoff)
        boilerplate_score = score_text(model, boilerplate_text)
        boilerplate_truncated_score = score_text(model, boilerplate_truncated)

        result = {
            "name": passage["name"],
            "total_words": total_words,
            "full_score": full_score,
            "truncated_mid_words": mid_cutoff,
            "truncated_mid_score": truncated_mid_score,
            "full_with_boilerplate_score": boilerplate_score,
            "truncated_from_boilerplate_score": boilerplate_truncated_score,
            "truncation_delta": truncated_mid_score - full_score,
            "boilerplate_removal_delta": boilerplate_truncated_score
            - boilerplate_score,
        }
        truncation_results.append(result)

        log.info(
            "%-20s full=%.4f  trunc_mid=%.4f  delta=%+.4f  "
            "w/boilerplate=%.4f  trunc_boilerplate=%.4f",
            passage["name"],
            full_score,
            truncated_mid_score,
            truncated_mid_score - full_score,
            boilerplate_score,
            boilerplate_truncated_score,
        )

        for pct in (25, 50, 75, 100):
            word_count = max(1, int(total_words * pct / 100))
            text_slice = truncate_words(passage["text"], word_count)
            length_score = score_text(model, text_slice)
            length_curve_results.append(
                {
                    "name": passage["name"],
                    "pct": pct,
                    "word_count": word_count,
                    "score": float(length_score),
                }
            )

    log.info("\n--- Length-quality curves ---\n")
    for passage in TRUNCATION_PASSAGES:
        entries = [r for r in length_curve_results if r["name"] == passage["name"]]
        scores_str = "  ".join(f"{e['pct']}%={e['score']:.4f}" for e in entries)
        log.info("%-20s %s", passage["name"], scores_str)

    avg_delta = sum(r["truncation_delta"] for r in truncation_results) / len(
        truncation_results
    )
    log.info("\nAvg truncation delta (trunc - full): %+.4f", avg_delta)
    if avg_delta > 0.01:
        log.info(
            "FINDING: Truncated text scores HIGHER on average (supports hypothesis B)"
        )
    elif avg_delta < -0.01:
        log.info("FINDING: Truncated text scores LOWER on average")
    else:
        log.info("FINDING: No meaningful difference from truncation")

    return truncation_results, length_curve_results


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        description="Phase 1: FastText quality model probe"
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Directory to write JSON results (optional)",
    )
    args = parser.parse_args()

    log.info("Loading FastText quality model...")
    from dolma.quality.fasttext import QualityFastTextClassifier

    classifier = QualityFastTextClassifier()
    log.info("Model loaded in %.2fs", classifier.load_time_seconds)

    raw_labels = list(classifier.model.get_labels())
    label_map = classifier.label_map
    log.info("Raw model labels: %s", raw_labels)
    log.info("Label map: %s", label_map.to_dict())

    result = ProbeResult(
        raw_model_labels=raw_labels,
        label_map=label_map.to_dict(),
    )

    result.part_a = run_part_a(classifier.model, label_map)
    result.part_b_truncation, result.part_b_length_curve = run_part_b(classifier.model)

    if args.output_dir is not None:
        args.output_dir.mkdir(parents=True, exist_ok=True)
        output_path = args.output_dir / "probe_fasttext_quality.json"
        with open(output_path, "w") as f:
            json.dump(asdict(result), f, indent=2)
        log.info("\nResults written to %s", output_path)


if __name__ == "__main__":
    main()