Function eval_compression_quality

test_compression.py:147–245

(
    logprob_rows: list[dict[str, torch.Tensor]],
    config: DistributionQuantizationConfig | LegacyLogitCompressionConfig,
    batch_size: int = 256,
    device: str = "cuda",
    missing_prob_handling: MissingProbabilityHandling = MissingProbabilityHandling.ZERO,
)

def eval_compression_quality(
    logprob_rows: list[dict[str, torch.Tensor]],
    config: DistributionQuantizationConfig | LegacyLogitCompressionConfig,
    batch_size: int = 256,
    device: str = "cuda",
    missing_prob_handling: MissingProbabilityHandling = MissingProbabilityHandling.ZERO,
) -> CompressionEvalResult:
    compressor = (
        LogprobCompressor(
            config=config
            if isinstance(config, DistributionQuantizationConfig)
            else None,
            legacy_config=(
                config if isinstance(config, LegacyLogitCompressionConfig) else None
            ),
        )
        if config != "bf16"
        else None
    )

    vocab_size = None
    bytes_per_token = None
    if config == "bf16":
        vocab_size = VOCAB_SIZE
        bytes_per_token = 2 * vocab_size
    elif isinstance(config, DistributionQuantizationConfig):
        vocab_size = config.d
        bytes_per_token = (config.total_bits() + 7) // 8
    elif isinstance(config, LegacyLogitCompressionConfig):
        vocab_size = config.vocab_size
        bytes_per_token = compressor.legacy_compressor.bytes_per_token()
    else:
        raise ValueError("Invalid config type")

    kld_fwd_list = []
    kld_bwd_list = []
    jsd_list = []
    mse_list = []
    for i_0 in tqdm.tqdm(range(0, len(logprob_rows), batch_size), desc="Evaluating"):
        i_1 = min(i_0 + batch_size, len(logprob_rows))
        batch = logprob_rows[i_0:i_1]
        logprobs = torch.stack([row["logprobs"] for row in batch], dim=0).to(device)
        mask = (torch.stack([row["attention_mask"] for row in batch], dim=0) > 0).to(
            device
        )
        logprobs = logprobs.float().masked_fill(mask.unsqueeze(-1), -1e6)

        if config != "bf16":
            batch_out = compressor.compress(logprobs)
            rt_ids, rt_logp = compressor.decompress_to_sparse(batch_out)
            del batch_out
            rt_logp_dense = densify(
                rt_ids,
                rt_logp.float(),
                vocab_size,
                missing=missing_prob_handling,
                renormalize=True,
                fill_value=-1e12,