perplexity.py

import json
import math

SENTINEL = -9999


def calculate_perplexity(
    api_response: dict,
    sentinel_strategy: str = "skip",  # "skip" | "clamp"
    clamp_value: float = -100.0,
) -> dict:
    """
    Calculate perplexity from an OpenAI chat completion response with logprobs.

    OpenAI occasionally emits -9999 as a sentinel for log-probs too small to
    represent.  Two strategies are supported:

      skip  — exclude sentinel tokens from the calculation entirely
      clamp — replace -9999 with `clamp_value` (default -100, i.e. p ≈ 3.7e-44)

    Returns a dict with perplexity and diagnostic counts.
    """
    if sentinel_strategy not in ("skip", "clamp"):
        raise ValueError("sentinel_strategy must be 'skip' or 'clamp'")

    choices = api_response.get("choices", [])
    if not choices:
        raise ValueError("No choices in response")

    logprobs_data = choices[0].get("logprobs")
    if not logprobs_data:
        raise ValueError("No logprobs in response. Request with logprobs=True.")

    raw = [t["logprob"] for t in logprobs_data["content"]]
    sentinels = [lp for lp in raw if lp == SENTINEL]

    if sentinel_strategy == "skip":
        filtered = [lp for lp in raw if lp != SENTINEL]
    else:  # clamp
        filtered = [clamp_value if lp == SENTINEL else lp for lp in raw]

    n = len(filtered)
    if n == 0:
        raise ValueError("No usable tokens after applying sentinel strategy.")

    perplexity = math.exp(-sum(filtered) / n)

    return {
        "perplexity": perplexity,
        "n_tokens_total": len(raw),
        "n_tokens_used": n,
        "n_sentinels": len(sentinels),
        "sentinel_strategy": sentinel_strategy,
    }


def summarize(api_response: dict, sentinel_strategy: str = "skip"):
    tokens = api_response["choices"][0]["logprobs"]["content"]

    print(f"{'Token':<20} {'LogProb':>10}  {'Prob':>12}  {'Note':>8}")
    print("-" * 56)
    for t in tokens[:15]:
        lp = t["logprob"]
        note = "SENTINEL" if lp == SENTINEL else ""
        prob = "—" if lp == SENTINEL else f"{math.exp(lp):.6f}"
        print(f"{repr(t['token']):<20} {lp:>10.4f}  {prob:>12}  {note:>8}")
    if len(tokens) > 15:
        print(f"  ... ({len(tokens) - 15} more tokens)")

    result = calculate_perplexity(api_response, sentinel_strategy=sentinel_strategy)

    print(f"\nTotal tokens    : {result['n_tokens_total']}")
    if result["n_sentinels"]:
        print(
            f"Sentinel (-9999): {result['n_sentinels']} token(s) [{sentinel_strategy}ped]"
        )
    print(f"Tokens used     : {result['n_tokens_used']}")
    print(f"Perplexity      : {result['perplexity']:.4f}")


if __name__ == "__main__":
    import sys

    path = sys.argv[1] if len(sys.argv) > 1 else "response.json"
    strategy = sys.argv[2] if len(sys.argv) > 2 else "skip"  # or "clamp"

    with open(path) as f:
        data = json.load(f)

    summarize(data, sentinel_strategy=strategy)