test_fidelity.py

"""
Does this Space read a PDF the way the corpus was read?

This is the test the whole design rests on. Part A's index and Part B's F1 of 0.945 were measured
on `skeleton_masked` and `payload_window` strings produced by the EDA notebook's extractor. If
`corpus_text.py` produces even slightly different text, those numbers stop describing this app and
become decoration.

So it pulls real PDFs out of the generation repo, runs them through `corpus_text.py`, and compares
the result **character by character** against the published parquet. Not "close enough" — identical.

The same argument applies to the family-naming model, for the same reason and with less margin: it
was fitted on seventeen numeric columns of that parquet, and `doc_features.py` recomputes them from
an uploaded file. A feature that drifts does not throw - it shifts the model's input distribution
and the model keeps answering confidently. So the numeric block is checked against the published
columns too, and the assembled feature row is checked for width against the fitted estimator.

    python test_fidelity.py [n_files]
"""

import sys

import numpy as np
import pandas as pd

import corpus_text
import doc_features

CORPUS = ("https://huggingface.co/datasets/Cyber-security-final-project/"
          "HARMLESS_Synthetic_Injected_PDFs_EDA/resolve/main/Datasets/"
          "synthetic_corpus_part2_clustered.parquet")
GENERATION_REPO = "Cyber-security-final-project/Generated_Injected_PDFs_HARMLESS"


def main(n=8):
    from huggingface_hub import hf_hub_download

    corpus = pd.read_parquet(CORPUS).set_index("file_id")

    # Injected and clean both, and not all from one family: masking and binary-stream handling
    # differ between them, and a test that only saw one would pass on a broken extractor.
    ids = list(corpus.index)
    sample = ids[::max(1, len(ids) // n)][:n]

    failures = 0
    for fid in sample:
        row = corpus.loc[fid]
        try:
            path = hf_hub_download(GENERATION_REPO, f"Output_PDFs/{fid}", repo_type="dataset")
        except Exception as e:
            print(f"  ? {fid}: could not fetch ({type(e).__name__})")
            continue

        data = open(path, "rb").read()
        skeleton, truncated, dropped = corpus_text.build_skeleton(data)
        masked = corpus_text.mask_leaks(skeleton)
        window = corpus_text.payload_window(skeleton)

        checks = {"skeleton_masked": (masked, row["skeleton_masked"]),
                  "payload_window": (window, row["payload_window"])}

        # The single-window path must also be what the triage returns first for a marked file,
        # or the model reads a different string here than the corpus was scored on.
        top = corpus_text.candidate_windows(skeleton, cover_all=False)[0]

        bad = [k for k, (got, want) in checks.items() if got != want]

        # The family model's document-shape block, against the columns it was fitted on. Compared
        # with a relative tolerance rather than exactly: the entropies and the compression ratio
        # are floats that went through a parquet round-trip, and demanding bit-equality of those
        # would fail on storage rather than on a real drift.
        stats = doc_features.describe(path, data, skeleton, truncated, dropped)
        drifted = [k for k in doc_features.NUMERIC
                   if k in row.index
                   and not np.isclose(float(stats[k]), float(row[k]), rtol=1e-4, atol=1e-6)]
        if drifted:
            bad += [f"feature {k}" for k in drifted]
            for k in drifted:
                print(f"      {k}: computed {stats[k]}, corpus has {row[k]}")

        if bad:
            failures += 1
            print(f"  x {fid}: differs in {', '.join(bad)}")
            for k in [b for b in bad if not b.startswith("feature ")]:
                got, want = checks[k]
                print(f"      {k}: got {len(got):,} chars, corpus has {len(want):,}")
        else:
            note = "head" if top["is_head"] else f"{len(top['families'])} family marker(s)"
            print(f"  . {fid}: identical  (triage top = {note}, "
                  f"{len(doc_features.NUMERIC)} features match)")

    print(f"\n{len(sample) - failures}/{len(sample)} identical to the published corpus")

    # Width of the assembled row against the fitted estimator. Cheap, and it is the check that
    # catches a retrain that added a feature block without anyone updating `build_features`.
    try:
        import family_model
        if family_model.AVAILABLE:
            row = family_model.build_features(
                skeleton, {"pred_family": "none", "pred_injected": 0, "parse_ok": 1},
                stats, np.zeros(len(family_model.CLASSES), dtype=np.float32))
            expected = family_model.load_model().n_features_in_
            ok = row.shape[0] == expected
            print(f"{'.' if ok else 'x'} feature row is {row.shape[0]} wide, "
                  f"the fitted model expects {expected}")
            failures += 0 if ok else 1
    except Exception as e:
        print(f"? family model not checked: {type(e).__name__}: {e}")

    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main(int(sys.argv[1]) if len(sys.argv) > 1 else 8))