exllamav3/eval/ppl.py

import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import argparse
from exllamav3.util.file import disk_lru_cache
from exllamav3.util.progress import ProgressBar
from exllamav3 import model_init, Tokenizer, Config
from datasets import load_dataset
from exllamav3.util.memory import free_mem
import torch
import torch.nn.functional as F
import math


@disk_lru_cache("get_dataset_text")
def get_dataset_text(spec: dict):
    assert spec["dataset"] == "wiki2", "Only wiki2 implemented atm"
    dataset_text = "\n\n".join(
        load_dataset("wikitext", "wikitext-2-raw-v1", split = "test")
        ["text"]
    )
    return dataset_text


def get_test_tokens(tokenizer, rows, eval_len = 2048, eval_stride = 2048):
    with ProgressBar("Tokenizing", rows) as pb:
        dataset_spec = { "dataset": "wiki2" }
        eval_tokens = tokenizer.encode(get_dataset_text(dataset_spec))
        num_tokens = eval_tokens.shape[-1]
        seqs = []
        for a in range(0, num_tokens - eval_len, eval_stride):
            b = a + eval_len
            seqs.append(eval_tokens[:, a:b])
            pb.update(len(seqs))
            if len(seqs) >= rows:
                break
    return torch.cat(seqs, dim = 0)[:, :]


@disk_lru_cache("_load_wikitext2_raw")
def _load_wikitext2_raw() -> str:
    """
    Get raw wikitext2 test split exactly as used by perplexity.c (datasets version differs slightly on whitespace)
    """

    import tempfile, zipfile, pathlib, urllib.request

    _WIKITEXT2_URL = "https://huggingface.co/datasets/ggml-org/ci/resolve/main/wikitext-2-raw-v1.zip"

    cache_dir = pathlib.Path(tempfile.gettempdir()) / "llama_cpp_ppl_wikitext2"
    cache_dir.mkdir(parents = True, exist_ok = True)

    raw_path = cache_dir / "wikitext-2-raw" / "wiki.test.raw"
    if not raw_path.exists():

        zip_path = cache_dir / "wikitext-2-raw-v1.zip"
        if not zip_path.exists():
            print(f"Downloading WikiText-2 raw to {zip_path} ...")
            urllib.request.urlretrieve(_WIKITEXT2_URL, str(zip_path))

        with zipfile.ZipFile(str(zip_path), "r") as zf:
            zf.extractall(str(cache_dir))

        zip_path.unlink(missing_ok = True)
        if not raw_path.exists():
            raise FileNotFoundError(f"Failed to extract to {raw_path}.")

    with open(raw_path, "r", encoding = "utf-8") as f:
        return f.read()


def get_test_tokens_gguf(tokenizer) -> list[int]:
    text = _load_wikitext2_raw()
    token_ids = tokenizer.encode(text, add_bos = False)
    return token_ids[0].tolist()


def eval_gguf(model, config, tokenizer, args, forward_fn):

    # Dataset
    tokens = get_test_tokens_gguf(tokenizer)
    bos_id = tokenizer.bos_token_id
    n_ctx = args.ctx_size

    if len(tokens) < 2 * n_ctx:
        raise ValueError(
            f"Need at least {2 * n_ctx} tokens for n_ctx={n_ctx}, "
            f"but the dataset tokenizes to only {len(tokens)} tokens."
        )

    # Chunking (non-overlapping)
    n_chunks = len(tokens) // n_ctx
    first = n_ctx // 2
    print(f"Perplexity: {len(tokens)} tokens, {n_chunks} chunks, n_ctx={n_ctx}, scoring positions [{first}, {n_ctx - 2}]")

    total_nll = 0.0
    total_nll2 = 0.0
    total_count = 0
    per_chunk_ppl: list[float] = []

    for i in range(n_chunks):
        start = i * n_ctx
        chunk_tokens = tokens[start : start + n_ctx]

        # Optionally overwrite first token in each chunk with BOS
        if bos_id is not None:
            chunk_tokens = [bos_id] + chunk_tokens[1:]

        input_ids = torch.tensor([chunk_tokens], dtype=torch.long)
        logits = forward_fn(model, input_ids)
        logits = logits[0, :, :tokenizer.actual_vocab_size].float()

        score_logits = logits[first : n_ctx - 1]  # [n_ctx - 1 - first, vocab]
        score_tokens = tokens[start + first + 1 : start + n_ctx]
        score_tokens = torch.tensor(score_tokens, dtype = torch.long, device = score_logits.device)

        # Compute per-token NLL via log-softmax (matches llama.cpp's log_softmax function: logit[tok] - max - log(sum_exp))
        log_probs = F.log_softmax(score_logits, dim=-1)
        token_nlls = -log_probs[torch.arange(len(score_tokens), device=logits.device), score_tokens]

        nll_sum = token_nlls.sum().item()
        nll2_sum = (token_nlls * token_nlls).sum().item()
        n_scored = len(score_tokens)
        total_nll += nll_sum
        total_nll2 += nll2_sum
        total_count += n_scored
        running_ppl = math.exp(total_nll / total_count)
        per_chunk_ppl.append(running_ppl)

        print(f"[{i + 1}]{running_ppl:.4f}", end=",", flush=True)

    print()

    # Aggregate
    mean_nll = total_nll / total_count
    ppl = math.exp(mean_nll)
    mean_nll2 = total_nll2 / total_count
    variance = mean_nll2 - mean_nll * mean_nll
    if variance > 0 and total_count > 1:
        nll_stderr = math.sqrt(variance / (total_count - 1))
    else:
        nll_stderr = 0.0
    ppl_stderr = nll_stderr * ppl

    print(f"Final estimate: PPL = {ppl:.4f} +/- {ppl_stderr:.5f}")


def eval_default(model, config, tokenizer, args, forward_fn):

    # Dataset
    eval_ids = get_test_tokens(tokenizer, args.rows, eval_len = args.length)
    vocab_size = tokenizer.actual_vocab_size

    # Test
    logprob_sum = 0.0
    logprob_count = 0
    with ProgressBar("Evaluating", args.rows) as pb:
        for row in range(eval_ids.shape[0]):
            pb.update(row)
            input_ids = eval_ids[row:row + 1, :]
            logits = forward_fn(model, input_ids)
            logits = logits[:, :-1, :vocab_size].float()
            logits += 1e-10
            log_probs = F.log_softmax(logits, dim = -1)
            del logits
            target_ids = input_ids[:, 1:].to(log_probs.device)
            del input_ids
            target_log_probs = log_probs.gather(-1, target_ids.unsqueeze(-1)).squeeze(-1)
            logprob_sum += target_log_probs.sum().item()
            logprob_count += target_ids.numel()
            del log_probs
            del target_log_probs
            del target_ids
            torch.cuda.empty_cache()
        pb.update(args.rows)
        mean_log_prob = logprob_sum / logprob_count
        perplexity = math.exp(-mean_log_prob)

    print(f" -- Evaluated: {eval_ids.shape[0]} rows of {eval_ids.shape[1]} tokens")
    print(f" -- Perplexity: {perplexity:.6f}")



@torch.inference_mode()
def main(args):

    if not args.hf:
        model, config, _, tokenizer = model_init.init(
            args,
            override_dynamic_seq_len = 2048,
            max_output_size = 2048,
            max_output_factor = 5,
        )
        def forward_fn_exl3(_model, _input_ids):
            return _model.forward(_input_ids, {"attn_mode": "flash_attn_nc"})
        forward_fn = forward_fn_exl3

        bpw_layer, bpw_head, vram_bits = model.get_storage_info()
        print(f" -- Model: {args.model_dir}")
        print(f" -- Bitrate: {bpw_layer:.2f} bpw / {bpw_head:.2f} bpw (head)")
    else:
        from transformers import AutoModelForCausalLM

        config = Config.from_directory(args.model_dir)
        tokenizer = Tokenizer.from_config(config)
        model = AutoModelForCausalLM.from_pretrained(
            args.model_dir,
            device_map = "auto",
            torch_dtype = torch.half if args.hf_tight else torch.float if args.hf_fp32 else None,
        )

        if args.hf_tight:
            free_mem()
            model.half()
            free_mem()
        if args.hf_fp32:
            free_mem()
            model.float()
            free_mem()

        def forward_fn_hf(_model, _input_ids):
            return _model.forward(_input_ids)["logits"]
        forward_fn = forward_fn_hf


    if not args.gguf:
        eval_default(model, config, tokenizer, args, forward_fn)
    else:
        eval_gguf(model, config, tokenizer, args, forward_fn)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    model_init.add_args(parser, cache = False)
    parser.add_argument("-r", "--rows", type = int, help = "Number of rows", default = 100)
    parser.add_argument("-l", "--length", type = int, help = "Length", default = 2048)
    parser.add_argument("-g", "--gguf", action = "store_true", help = "Use GGUF-equivalent eval logic (ignores -r and -l)")
    parser.add_argument("-c", "--ctx-size", type = int, help = "For GGUF-equiv.: size of the prompt context (default: 512)", default = 512)
    parser.add_argument("-hf", "--hf", action = "store_true", help = "Use Transformers as backend (-m must be HF model)")
    parser.add_argument("-hf_t", "--hf_tight", action = "store_true", help = "For Transformers: Force FP16 dtype to save memory")
    parser.add_argument("-hf_fp32", "--hf_fp32", action = "store_true", help = "For Transformers: Force FP32 dtype")

    _args = parser.parse_args()
    main(_args)