src/ppl_eval.py

#!/usr/bin/env python3
"""Perplexity evaluation for causal LMs on HF datasets or provided text."""

import argparse
import json
import math
import os
from typing import Dict, Iterable, List, Optional

import torch

try:
    from datasets import load_dataset
except Exception:  # pragma: no cover - optional dependency
    load_dataset = None

try:
    from transformers import AutoModelForCausalLM, AutoTokenizer
except Exception as exc:  # pragma: no cover - fail early with clear error
    raise SystemExit("transformers is required: pip install transformers") from exc

try:
    from tqdm import tqdm
except Exception:  # pragma: no cover - optional dependency
    tqdm = None


def _tqdm_enabled() -> bool:
    value = os.environ.get("DISABLE_TQDM", os.environ.get("TQDM_DISABLE", "0"))
    return value.strip().lower() not in {"1", "true", "yes", "on"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compute perplexity for a causal LM on one or more datasets."
    )
    parser.add_argument("--model", required=True, help="HF model id or local path")
    parser.add_argument(
        "--dataset",
        action="append",
        default=[],
        help="HF dataset name (repeatable).",
    )
    parser.add_argument(
        "--dataset_config",
        action="append",
        default=[],
        help="Optional dataset config (repeatable or single shared config).",
    )
    parser.add_argument(
        "--dataset_split",
        default="test",
        help="Dataset split to use (default: test)",
    )
    parser.add_argument(
        "--dataset_text_field",
        default=None,
        help="Text field in dataset (default: auto-detect, applies to all datasets)",
    )
    parser.add_argument(
        "--text",
        action="append",
        default=[],
        help="Inline text samples (can pass multiple)",
    )
    parser.add_argument(
        "--text_file",
        default=None,
        help="Path to a text file for evaluation data",
    )
    parser.add_argument(
        "--num_samples",
        type=int,
        default=0,
        help="Number of token sequences to use per dataset (0 = all)",
    )
    parser.add_argument(
        "--seq_len", type=int, default=2048, help="Sequence length"
    )
    parser.add_argument(
        "--batch_size", type=int, default=2, help="Batch size"
    )
    parser.add_argument(
        "--max_batches",
        type=int,
        default=None,
        help="Optional max number of batches to evaluate per dataset",
    )
    parser.add_argument(
        "--model_family",
        type=str,
        choices=["auto", "llama", "qwen"],
        default="auto",
        help="Model family for BOS handling",
    )
    parser.add_argument(
        "--add_bos",
        type=str,
        choices=["auto", "always", "never"],
        default="auto",
        help="Whether to prepend BOS to each sample",
    )
    parser.add_argument(
        "--device",
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device for model + compute",
    )
    parser.add_argument(
        "--dtype",
        default="auto",
        choices=["auto", "float32", "float16", "bfloat16"],
        help="Model dtype",
    )
    parser.add_argument(
        "--seed", type=int, default=0, help="Random seed for shuffling"
    )
    parser.add_argument(
        "--shuffle",
        action="store_true",
        help="Shuffle dataset before sampling",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=0,
        help="DataLoader workers",
    )
    parser.add_argument(
        "--cache_dir",
        default=None,
        help="Optional datasets cache directory",
    )
    parser.add_argument(
        "--trust_remote_code",
        action="store_true",
        help="Allow custom model code from hub",
    )
    parser.add_argument(
        "--output",
        default=None,
        help="Optional JSON output path",
    )
    return parser.parse_args()


def _normalize_config(config: Optional[str]) -> Optional[str]:
    if config is None:
        return None
    if config.strip().lower() in {"none", "null", "-"}:
        return None
    return config


def _expand_dataset_configs(
    datasets: List[str], configs: List[str]
) -> List[Optional[str]]:
    if not configs:
        return [None] * len(datasets)
    if len(configs) == 1 and len(datasets) > 1:
        return [_normalize_config(configs[0])] * len(datasets)
    if len(configs) != len(datasets):
        raise SystemExit(
            "Provide zero, one, or matching-count --dataset_config values."
        )
    return [_normalize_config(cfg) for cfg in configs]


def guess_text_field(dataset) -> str:
    if hasattr(dataset, "column_names") and dataset.column_names:
        if "text" in dataset.column_names:
            return "text"
        return dataset.column_names[0]
    if hasattr(dataset, "features"):
        names = list(dataset.features.keys())
        if "text" in names:
            return "text"
        if names:
            return names[0]
    return "text"


def _infer_model_family(model) -> str:
    model_type = str(getattr(getattr(model, "config", None), "model_type", "")).lower()
    architectures = getattr(getattr(model, "config", None), "architectures", [])
    arch_lower = " ".join(str(name).lower() for name in architectures)
    if "qwen" in model_type or "qwen" in arch_lower:
        return "qwen"
    if "llama" in model_type or "llama" in arch_lower:
        return "llama"
    return "unknown"


def _resolve_add_bos(setting: str, model_family: str, tokenizer) -> bool:
    if setting == "always":
        return True
    if setting == "never":
        return False
    if model_family == "llama":
        return True
    if model_family == "qwen":
        return False
    if hasattr(tokenizer, "add_bos_token"):
        return bool(getattr(tokenizer, "add_bos_token"))
    init_kwargs = getattr(tokenizer, "init_kwargs", None)
    if isinstance(init_kwargs, dict) and "add_bos_token" in init_kwargs:
        return bool(init_kwargs["add_bos_token"])
    return False


def build_token_chunks(
    texts: Iterable[str],
    tokenizer,
    seq_len: int,
    num_samples: int,
    add_bos: bool = False,
) -> List[torch.Tensor]:
    chunks: List[torch.Tensor] = []
    buffer: List[int] = []
    for text in texts:
        ids = tokenizer.encode(text, add_special_tokens=False)
        if add_bos and tokenizer.bos_token_id is not None:
            ids = [tokenizer.bos_token_id] + ids
        if not ids:
            continue
        buffer.extend(ids)
        while len(buffer) >= seq_len and len(chunks) < num_samples:
            chunk = buffer[:seq_len]
            buffer = buffer[seq_len:]
            chunks.append(torch.tensor(chunk, dtype=torch.long))
        if len(chunks) >= num_samples:
            break
    return chunks


def get_dtype(dtype: str):
    if dtype == "auto":
        return None
    if dtype == "float16":
        return torch.float16
    if dtype == "bfloat16":
        return torch.bfloat16
    return torch.float32


def compute_ppl(model, dataloader, device: str, max_batches: Optional[int]) -> float:
    model.eval()
    nll_sum = 0.0
    token_count = 0
    iterator = dataloader
    if tqdm is not None and _tqdm_enabled():
        iterator = tqdm(dataloader, desc="PPL", unit="batch")
    with torch.no_grad():
        for step, batch in enumerate(iterator):
            if isinstance(batch, dict):
                input_ids = batch["input_ids"].to(device)
            else:
                input_ids = batch[0].to(device)
            outputs = model(input_ids=input_ids)
            logits = outputs.logits
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = input_ids[:, 1:].contiguous()
            loss = torch.nn.functional.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                reduction="sum",
            )
            nll_sum += float(loss.item())
            token_count += shift_labels.numel()
            if max_batches is not None and step + 1 >= max_batches:
                break

    if token_count == 0:
        raise RuntimeError("No tokens processed; check evaluation inputs.")

    return math.exp(nll_sum / token_count)


def _load_lm_dataset(
    tokenizer,
    dataset_name: str,
    config: Optional[str],
    split: str,
    text_field: Optional[str],
    seq_len: int,
    add_bos: bool,
    cache_dir: Optional[str],
):
    dataset = load_dataset(
        dataset_name,
        config,
        split=split,
        trust_remote_code=True,
        cache_dir=cache_dir,
    )

    field = text_field or guess_text_field(dataset)

    def is_valid_text(example) -> bool:
        value = example.get(field)
        return isinstance(value, str) and value.strip() != ""

    dataset = dataset.filter(is_valid_text, desc=f"filter-{dataset_name}")

    def tokenize_fn(examples):
        tokenized = tokenizer(
            examples[field],
            add_special_tokens=False,
            return_attention_mask=False,
        )
        if add_bos and tokenizer.bos_token_id is not None:
            tokenized["input_ids"] = [
                [tokenizer.bos_token_id] + ids for ids in tokenized["input_ids"]
            ]
        return tokenized

    tokenized = dataset.map(
        tokenize_fn,
        batched=True,
        remove_columns=dataset.column_names,
        desc=f"tokenize-{dataset_name}",
    )

    def group_texts(examples):
        concatenated = []
        for ids in examples["input_ids"]:
            concatenated.extend(ids)
        total_length = (len(concatenated) // seq_len) * seq_len
        if total_length == 0:
            return {"input_ids": []}
        return {
            "input_ids": [
                concatenated[i : i + seq_len] for i in range(0, total_length, seq_len)
            ]
        }

    lm_dataset = tokenized.map(
        group_texts,
        batched=True,
        batch_size=1000,
        remove_columns=tokenized.column_names,
        desc=f"group-{dataset_name}",
    )
    lm_dataset.set_format(type="torch", columns=["input_ids"])
    return lm_dataset


def prepare_ppl_dataloaders(
    tokenizer,
    datasets: List[str],
    configs: List[Optional[str]],
    split: str,
    text_field: Optional[str],
    num_samples: int,
    seq_len: int,
    batch_size: int,
    seed: int,
    shuffle: bool,
    model_family: str = "auto",
    add_bos: str = "auto",
    cache_dir: Optional[str] = None,
    num_workers: int = 0,
    model=None,
) -> Dict[str, torch.utils.data.DataLoader]:
    if load_dataset is None:
        raise SystemExit("datasets is required for dataset evaluation")

    resolved_family = model_family
    if resolved_family == "auto":
        if model is None:
            raise SystemExit("model is required when model_family is 'auto'")
        resolved_family = _infer_model_family(model)
    use_bos = _resolve_add_bos(add_bos, resolved_family, tokenizer)
    if use_bos and tokenizer.bos_token_id is None:
        use_bos = False

    dataloaders: Dict[str, torch.utils.data.DataLoader] = {}
    for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
        lm_dataset = _load_lm_dataset(
            tokenizer=tokenizer,
            dataset_name=dataset_name,
            config=config,
            split=split,
            text_field=text_field,
            seq_len=seq_len,
            add_bos=use_bos,
            cache_dir=cache_dir,
        )
        if shuffle:
            try:
                lm_dataset = lm_dataset.shuffle(seed=seed + idx)
            except Exception:
                pass
        if num_samples and hasattr(lm_dataset, "__len__"):
            limit = min(num_samples, len(lm_dataset))
            lm_dataset = lm_dataset.select(range(limit))

        data_loader = torch.utils.data.DataLoader(
            lm_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
        )
        label = dataset_name if config is None else f"{dataset_name}:{config}"
        dataloaders[label] = data_loader

    return dataloaders


def evaluate_ppl_dataloaders(
    model,
    dataloaders: Dict[str, torch.utils.data.DataLoader],
    device: str,
    max_batches: Optional[int] = None,
) -> Dict[str, float]:
    results: Dict[str, float] = {}
    for label, data_loader in dataloaders.items():
        ppl = compute_ppl(model, data_loader, device, max_batches=max_batches)
        results[label] = ppl
    return results


def evaluate_ppl_datasets(
    model,
    tokenizer,
    datasets: List[str],
    configs: List[Optional[str]],
    split: str,
    text_field: Optional[str],
    num_samples: int,
    seq_len: int,
    batch_size: int,
    device: str,
    seed: int,
    shuffle: bool,
    model_family: str = "auto",
    add_bos: str = "auto",
    max_batches: Optional[int] = None,
    cache_dir: Optional[str] = None,
    num_workers: int = 0,
) -> Dict[str, float]:
    if load_dataset is None:
        raise SystemExit("datasets is required for dataset evaluation")

    resolved_family = model_family
    if resolved_family == "auto":
        resolved_family = _infer_model_family(model)
    use_bos = _resolve_add_bos(add_bos, resolved_family, tokenizer)
    if use_bos and tokenizer.bos_token_id is None:
        use_bos = False

    results: Dict[str, float] = {}
    for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
        lm_dataset = _load_lm_dataset(
            tokenizer=tokenizer,
            dataset_name=dataset_name,
            config=config,
            split=split,
            text_field=text_field,
            seq_len=seq_len,
            add_bos=use_bos,
            cache_dir=cache_dir,
        )
        if shuffle:
            try:
                lm_dataset = lm_dataset.shuffle(seed=seed + idx)
            except Exception:
                pass
        if num_samples and hasattr(lm_dataset, "__len__"):
            limit = min(num_samples, len(lm_dataset))
            lm_dataset = lm_dataset.select(range(limit))

        data_loader = torch.utils.data.DataLoader(
            lm_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
        )
        label = dataset_name if config is None else f"{dataset_name}:{config}"
        ppl = compute_ppl(model, data_loader, device, max_batches=max_batches)
        results[label] = ppl
    return results


def main() -> None:
    args = parse_args()
    torch.manual_seed(args.seed)

    dtype = get_dtype(args.dtype)
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        torch_dtype=dtype,
        trust_remote_code=args.trust_remote_code,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        args.model, trust_remote_code=args.trust_remote_code
    )
    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token

    model.to(args.device)

    results: Dict[str, float] = {}
    resolved_family = args.model_family
    if resolved_family == "auto":
        resolved_family = _infer_model_family(model)
    use_bos = _resolve_add_bos(args.add_bos, resolved_family, tokenizer)
    if use_bos and tokenizer.bos_token_id is None:
        use_bos = False

    if args.dataset:
        datasets = list(args.dataset)
        configs = _expand_dataset_configs(datasets, list(args.dataset_config))
        results.update(
            evaluate_ppl_datasets(
                model,
                tokenizer,
                datasets=datasets,
                configs=configs,
                split=args.dataset_split,
                text_field=args.dataset_text_field,
                num_samples=args.num_samples,
                seq_len=args.seq_len,
                batch_size=args.batch_size,
                device=args.device,
                seed=args.seed,
                shuffle=args.shuffle,
                model_family=resolved_family,
                add_bos="always" if use_bos else "never",
                max_batches=args.max_batches,
                cache_dir=args.cache_dir,
                num_workers=args.num_workers,
            )
        )

    if args.text_file or args.text:
        custom_texts: List[str] = []
        if args.text_file:
            with open(args.text_file, "r", encoding="utf-8") as handle:
                custom_texts.extend([line.strip() for line in handle if line.strip()])
        if args.text:
            custom_texts.extend([t for t in args.text if t])
        if custom_texts:
            chunks = build_token_chunks(
                custom_texts,
                tokenizer,
                args.seq_len,
                args.num_samples if args.num_samples > 0 else 1_000_000,
                add_bos=use_bos,
            )
            if not chunks:
                raise SystemExit(
                    "Not enough custom text to build token sequences. "
                    "Provide more --text/--text_file content or reduce --seq_len."
                )
            dataset = torch.utils.data.TensorDataset(torch.stack(chunks))
            dataloader = torch.utils.data.DataLoader(
                dataset, batch_size=args.batch_size, shuffle=False
            )
            results["custom"] = compute_ppl(
                model, dataloader, args.device, max_batches=args.max_batches
            )

    if not results:
        raise SystemExit("Provide --dataset and/or --text/--text_file for evaluation")

    print("Perplexity results:")
    for name, ppl in results.items():
        print(f"{name}: {ppl:.4f}")

    if args.output:
        with open(args.output, "w", encoding="utf-8") as handle:
            json.dump({"model": args.model, "results": results}, handle, indent=2)


if __name__ == "__main__":
    main()