#!/usr/bin/env python3
"""How many BPE tokens is each passage, and which ones did the scorer truncate?

`score.py` tokenizes with `truncation=True, max_length=512`. Passages run 240-430 words, which is
roughly 320-570 tokens, so some fraction of them are scored on less text than they contain. That
is a confound nobody would see in the output: a truncated passage is silently a *different, shorter*
passage as far as the detector is concerned, and if truncation correlates with the flag rate then
the length analysis is really measuring truncation.

Both detectors are roberta-base, so both should use the same BPE vocabulary — this computes the
count under each tokenizer and asserts they agree rather than assuming it. Tokenizers load in a
few MB, with no model weights, so this runs safely beside a live scan.

Writes tokens.jsonl: {"n", "tok", "trunc"} for the adversarial corpus.
"""
import json, os, sys

HERE = os.path.dirname(os.path.abspath(__file__))
MAX = 512


def main():
    from transformers import AutoTokenizer
    a = AutoTokenizer.from_pretrained("Hello-SimpleAI/chatgpt-detector-roberta")
    b = AutoTokenizer.from_pretrained("openai-community/roberta-base-openai-detector")

    rows = [json.loads(l) for l in open(os.path.join(HERE, "passages.jsonl"))]
    disagree = 0
    trunc = 0
    with open(os.path.join(HERE, "tokens.jsonl"), "w") as fh:
        for i, r in enumerate(rows):
            na = len(a(r["text"])["input_ids"])
            nb = len(b(r["text"])["input_ids"])
            if na != nb:
                disagree += 1
            t = na > MAX
            trunc += t
            fh.write(json.dumps({"n": r["n"], "tok": na, "trunc": bool(t)}) + "\n")
            if i % 2000 == 0:
                print("  %d/%d" % (i, len(rows)), file=sys.stderr, flush=True)

    print("passages: %d   truncated at %d tokens: %d (%.1f%%)"
          % (len(rows), MAX, trunc, 100 * trunc / len(rows)), file=sys.stderr)
    print("tokenizers disagreeing on length: %d %s"
          % (disagree, "(identical vocab, as expected)" if not disagree
             else "** the two detectors do NOT see the same tokens **"), file=sys.stderr)
    print("done", file=sys.stderr, flush=True)


if __name__ == "__main__":
    main()