feat: rewrite evaluate_ppl with token-level loss and multi-file support
- Support multiple input files, glob patterns, and directory input - Add --token_level flag: per-record token_ids + log_probs JSONL output - Add --max_samples for random subsampling per file - LossAccumulator: streaming mode (histogram-based percentiles, low memory) vs exact mode (full token list) - Token type analysis (ascii/cjk/non_ascii/special) when token_level=True - Fix token_ids/log_probs alignment (shift offset) - Cache frozenset(stop_ids) outside loop for performance - Aggregate stats: mean/median/ppl/p50/p90/p95/p99 - Summary JSON with all datasets in one file
This commit is contained in:
parent
cd14d53707
commit
b12b24eadc
|
|
@ -1,5 +1,9 @@
|
|||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
|
@ -9,95 +13,395 @@ from astrai.model import AutoModel
|
|||
from astrai.tokenize import AutoTokenizer
|
||||
|
||||
|
||||
def _collect_input_files(input_path: str) -> List[str]:
|
||||
"""Resolve *input_path* to a list of JSONL/JSON files."""
|
||||
if os.path.isdir(input_path):
|
||||
files = []
|
||||
for ext in ("*.jsonl", "*.json"):
|
||||
files.extend(
|
||||
sorted(glob.glob(os.path.join(input_path, "**", ext), recursive=True))
|
||||
)
|
||||
return files
|
||||
return sorted(glob.glob(input_path))
|
||||
|
||||
|
||||
def _load_items(filepath: str) -> List[dict]:
|
||||
"""Load JSONL or JSON (array / single dict) into a list of dicts."""
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
if filepath.lower().endswith(".json"):
|
||||
data = json.load(f)
|
||||
if isinstance(data, dict):
|
||||
return [data]
|
||||
return data
|
||||
return [json.loads(line) for line in f if line.strip()]
|
||||
|
||||
|
||||
def _encode_batch(
|
||||
tokenizer: AutoTokenizer, texts: List[str], max_length: int
|
||||
) -> Tuple[List[List[int]], List[List[int]]]:
|
||||
"""Encode *texts* and return (token_ids, attention_masks).
|
||||
|
||||
Each sequence is left-aligned and padded to the batch max length.
|
||||
"""
|
||||
encoded = [tokenizer.encode(t)[:max_length] for t in texts]
|
||||
if not encoded:
|
||||
return [], []
|
||||
max_len = max(len(seq) for seq in encoded)
|
||||
padded_ids = []
|
||||
masks = []
|
||||
for seq in encoded:
|
||||
pad_len = max_len - len(seq)
|
||||
padded_ids.append(seq + [tokenizer.pad_id] * pad_len)
|
||||
masks.append([1] * len(seq) + [0] * pad_len)
|
||||
return padded_ids, masks
|
||||
|
||||
|
||||
def _compute_batch(
|
||||
model,
|
||||
input_ids: torch.Tensor,
|
||||
attention_mask: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Forward pass and return (log_probs, valid_mask) of shape [B, S-1].
|
||||
|
||||
log_probs[i, j] = log P(token j+1 | tokens 0..j)
|
||||
"""
|
||||
output = model(input_ids, input_mask=attention_mask)
|
||||
logits = output["logits"][:, :-1, :] # [B, S-1, V]
|
||||
targets = input_ids[:, 1:] # [B, S-1]
|
||||
valid = attention_mask[:, 1:].float() # [B, S-1]
|
||||
|
||||
log_probs = F.log_softmax(logits.float(), dim=-1) # [B, S-1, V]
|
||||
token_log_probs = log_probs.gather(2, targets.unsqueeze(-1)).squeeze(-1) # [B, S-1]
|
||||
|
||||
return token_log_probs, valid
|
||||
|
||||
|
||||
def _token_type(token_id: int, stop_ids: frozenset, decode_fn) -> str:
|
||||
"""Classify a token into a coarse type for analysis.
|
||||
|
||||
*stop_ids* is a pre-built set of special token IDs.
|
||||
*decode_fn* is ``tokenizer.decode`` (or a wrapper) for single-token
|
||||
decoding.
|
||||
"""
|
||||
if token_id in stop_ids:
|
||||
return "special"
|
||||
decoded = decode_fn([token_id], skip_special_tokens=True)
|
||||
if any("\u4e00" <= ch <= "\u9fff" for ch in decoded):
|
||||
return "cjk"
|
||||
if any(ord(ch) > 127 for ch in decoded):
|
||||
return "non_ascii"
|
||||
return "ascii"
|
||||
|
||||
|
||||
def _percentiles(values: List[float]) -> Dict[str, float]:
|
||||
"""Compute common percentiles from a list of floats.
|
||||
|
||||
Uses linear interpolation between closest ranks (same convention
|
||||
as NumPy's default).
|
||||
"""
|
||||
if not values:
|
||||
return {}
|
||||
sorted_vals = sorted(values)
|
||||
n = len(sorted_vals)
|
||||
|
||||
def _pct(p: float) -> float:
|
||||
if n == 1:
|
||||
return sorted_vals[0]
|
||||
k = p * (n - 1)
|
||||
f = int(k)
|
||||
c = min(f + 1, n - 1)
|
||||
return sorted_vals[f] + (sorted_vals[c] - sorted_vals[f]) * (k - f)
|
||||
|
||||
return {
|
||||
"p50": _pct(0.50),
|
||||
"p90": _pct(0.90),
|
||||
"p95": _pct(0.95),
|
||||
"p99": _pct(0.99),
|
||||
}
|
||||
|
||||
|
||||
class LossAccumulator:
|
||||
"""Accumulate per-token losses with optional streaming mode.
|
||||
|
||||
When *stream* is True (token_level=False), losses are not kept
|
||||
in memory individually — only a running sum/count and a histogram
|
||||
(for approximate percentiles) are maintained. When *stream* is
|
||||
False, all losses are retained for exact statistics and per-record
|
||||
output.
|
||||
"""
|
||||
|
||||
_HIST_BINS = 1000
|
||||
_HIST_MAX = 20.0 # clamp losses above this for histogram
|
||||
|
||||
def __init__(self, stream: bool):
|
||||
self.stream = stream
|
||||
self.losses: List[float] = [] if not stream else []
|
||||
self.total: float = 0.0
|
||||
self.count: int = 0
|
||||
self.hist = torch.zeros(self._HIST_BINS, dtype=torch.long)
|
||||
# per-type losses (only populated when not streaming)
|
||||
self.by_type: Dict[str, List[float]] = {}
|
||||
self.type_total: Dict[str, float] = {}
|
||||
self.type_count: Dict[str, int] = {}
|
||||
|
||||
def add(self, losses: List[float]):
|
||||
self.total += sum(losses)
|
||||
self.count += len(losses)
|
||||
if self.stream:
|
||||
clamped = [min(max(l, 0.0), self._HIST_MAX) for l in losses]
|
||||
idx = torch.tensor(clamped) / self._HIST_MAX * (self._HIST_BINS - 1)
|
||||
self.hist += torch.bincount(
|
||||
idx.long().clamp(0, self._HIST_BINS - 1),
|
||||
minlength=self._HIST_BINS,
|
||||
)
|
||||
else:
|
||||
self.losses.extend(losses)
|
||||
|
||||
def add_typed(self, ttype: str, losses: List[float]):
|
||||
if not self.stream:
|
||||
self.by_type.setdefault(ttype, []).extend(losses)
|
||||
self.type_total[ttype] = self.type_total.get(ttype, 0.0) + sum(losses)
|
||||
self.type_count[ttype] = self.type_count.get(ttype, 0) + len(losses)
|
||||
|
||||
def stats(self) -> Dict:
|
||||
result: Dict = {}
|
||||
if self.count == 0:
|
||||
return result
|
||||
mean_loss = self.total / self.count
|
||||
result["overall"] = {
|
||||
"num_tokens": self.count,
|
||||
"mean_loss": mean_loss,
|
||||
"ppl": float(torch.exp(torch.tensor(mean_loss))),
|
||||
}
|
||||
if self.stream:
|
||||
result["overall"].update(self._hist_percentiles())
|
||||
else:
|
||||
result["overall"]["median_loss"] = statistics.median(self.losses)
|
||||
result["overall"].update(_percentiles(self.losses))
|
||||
|
||||
if self.type_count:
|
||||
result["by_token_type"] = {}
|
||||
for ttype in sorted(self.type_count.keys()):
|
||||
cnt = self.type_count[ttype]
|
||||
tmean = self.type_total[ttype] / cnt
|
||||
entry: Dict = {
|
||||
"num_tokens": cnt,
|
||||
"mean_loss": tmean,
|
||||
"ppl": float(torch.exp(torch.tensor(tmean))),
|
||||
}
|
||||
if not self.stream and ttype in self.by_type:
|
||||
entry["median_loss"] = statistics.median(self.by_type[ttype])
|
||||
entry.update(_percentiles(self.by_type[ttype]))
|
||||
result["by_token_type"][ttype] = entry
|
||||
return result
|
||||
|
||||
def _hist_percentiles(self) -> Dict[str, float]:
|
||||
"""Approximate percentiles from the histogram."""
|
||||
total = self.hist.sum().item()
|
||||
if total == 0:
|
||||
return {}
|
||||
cum = torch.cumsum(self.hist.float(), dim=0)
|
||||
result = {}
|
||||
for label, p in [("p50", 0.5), ("p90", 0.9), ("p95", 0.95), ("p99", 0.99)]:
|
||||
target = p * total
|
||||
idx = int(torch.searchsorted(cum, target).item())
|
||||
idx = min(idx, self._HIST_BINS - 1)
|
||||
result[label] = (idx + 0.5) / self._HIST_BINS * self._HIST_MAX
|
||||
return result
|
||||
|
||||
|
||||
def process_file(
|
||||
param_path: str, input_file: str, output_file: str, batch_size: int, text_key: str
|
||||
model,
|
||||
tokenizer: AutoTokenizer,
|
||||
items: List[dict],
|
||||
text_key: str,
|
||||
batch_size: int,
|
||||
max_length: int,
|
||||
token_level: bool,
|
||||
max_samples: Optional[int],
|
||||
output_file: Optional[str],
|
||||
label: str,
|
||||
) -> Dict:
|
||||
"""Evaluate a single dataset (list of items), return summary stats.
|
||||
|
||||
If *token_level* is True and *output_file* is set, per-record token_ids
|
||||
and log_probs are written as JSONL alongside the summary.
|
||||
"""
|
||||
if max_samples and len(items) > max_samples:
|
||||
import random
|
||||
|
||||
items = random.sample(items, max_samples)
|
||||
|
||||
texts = [item[text_key] for item in items if text_key in item]
|
||||
print(f" [{label}] {len(texts)} samples, text_key='{text_key}'")
|
||||
|
||||
acc = LossAccumulator(stream=not token_level)
|
||||
per_sample: List[dict] = []
|
||||
|
||||
if token_level:
|
||||
stop_ids = frozenset(tokenizer.stop_ids)
|
||||
decode_fn = tokenizer.decode
|
||||
|
||||
num_batches = (len(texts) + batch_size - 1) // batch_size
|
||||
for i in tqdm.tqdm(
|
||||
range(0, len(texts), batch_size),
|
||||
total=num_batches,
|
||||
desc=f" {label}",
|
||||
leave=False,
|
||||
):
|
||||
batch_texts = texts[i : i + batch_size]
|
||||
padded_ids, masks = _encode_batch(tokenizer, batch_texts, max_length)
|
||||
|
||||
input_ids = torch.tensor(padded_ids, device="cuda", dtype=torch.long)
|
||||
attention_mask = torch.tensor(masks, device="cuda", dtype=torch.bool)
|
||||
|
||||
token_log_probs, valid = _compute_batch(model, input_ids, attention_mask)
|
||||
|
||||
for b in range(len(batch_texts)):
|
||||
seq_len = int(valid[b].sum().item())
|
||||
lps = token_log_probs[b, :seq_len].tolist()
|
||||
losses = [-lp for lp in lps]
|
||||
acc.add(losses)
|
||||
|
||||
if token_level:
|
||||
# log_probs correspond to positions 1..seq_len (predicted
|
||||
# from position 0..seq_len-1), so token_ids must skip BOS
|
||||
# at position 0 to stay aligned with log_probs.
|
||||
ids = padded_ids[b][1 : seq_len + 1]
|
||||
per_sample.append(
|
||||
{
|
||||
"text": batch_texts[b][:200],
|
||||
"token_ids": ids,
|
||||
"log_probs": [round(lp, 4) for lp in lps],
|
||||
"ppl": float(torch.exp(torch.tensor(statistics.mean(losses))))
|
||||
if losses
|
||||
else None,
|
||||
}
|
||||
)
|
||||
typed_losses: Dict[str, List[float]] = {}
|
||||
for tid, loss in zip(ids, losses):
|
||||
ttype = _token_type(tid, stop_ids, decode_fn)
|
||||
typed_losses.setdefault(ttype, []).append(loss)
|
||||
for ttype, tl in typed_losses.items():
|
||||
acc.add_typed(ttype, tl)
|
||||
|
||||
stats = acc.stats()
|
||||
|
||||
if token_level and output_file:
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
for item in per_sample:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
|
||||
return stats
|
||||
|
||||
|
||||
def print_stats(label: str, stats: Dict):
|
||||
"""Pretty-print summary statistics."""
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" {label}")
|
||||
print(f"{'=' * 60}")
|
||||
ov = stats.get("overall", {})
|
||||
if ov:
|
||||
print(f" tokens: {ov['num_tokens']:,}")
|
||||
print(f" mean loss: {ov['mean_loss']:.4f}")
|
||||
if "median_loss" in ov:
|
||||
print(f" median loss: {ov['median_loss']:.4f}")
|
||||
print(f" ppl: {ov['ppl']:.2f}")
|
||||
if "p50" in ov:
|
||||
print(
|
||||
f" p50/p90/p95/p99: "
|
||||
f"{ov['p50']:.2f} / {ov['p90']:.2f} / {ov['p95']:.2f} / {ov['p99']:.2f}"
|
||||
)
|
||||
by_type = stats.get("by_token_type", {})
|
||||
if by_type:
|
||||
print(f"\n by token type:")
|
||||
print(f" {'type':<12} {'count':>8} {'mean_loss':>10} {'ppl':>8}")
|
||||
print(f" {'-' * 12} {'-' * 8} {'-' * 10} {'-' * 8}")
|
||||
for ttype, s in by_type.items():
|
||||
print(
|
||||
f" {ttype:<12} {s['num_tokens']:>8,} "
|
||||
f"{s['mean_loss']:>10.4f} {s['ppl']:>8.2f}"
|
||||
)
|
||||
|
||||
|
||||
def main(
|
||||
param_path: str,
|
||||
input_path: str,
|
||||
output_dir: str,
|
||||
text_key: str,
|
||||
batch_size: int,
|
||||
max_length: int,
|
||||
token_level: bool,
|
||||
max_samples: Optional[int],
|
||||
):
|
||||
# Load model and tokenizer
|
||||
print(f"Loading model from {param_path} ...")
|
||||
model = AutoModel.from_pretrained(param_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(param_path)
|
||||
model.to(device="cuda", dtype=torch.bfloat16)
|
||||
model.eval()
|
||||
|
||||
with open(input_file, "r", encoding="utf-8") as f:
|
||||
input_data = [json.loads(line) for line in f]
|
||||
input_files = _collect_input_files(input_path)
|
||||
if not input_files:
|
||||
print(f"No input files found at {input_path}")
|
||||
return
|
||||
|
||||
texts = [item[text_key] for item in input_data]
|
||||
print(f"Found {len(input_files)} file(s) to evaluate")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Encode all texts
|
||||
print(f"Encoding {len(texts)} texts...")
|
||||
encoded_texts = [tokenizer.encode(text) for text in texts]
|
||||
all_stats = {}
|
||||
for filepath in input_files:
|
||||
label = os.path.splitext(os.path.basename(filepath))[0]
|
||||
items = _load_items(filepath)
|
||||
if not items:
|
||||
print(f" [{label}] empty, skipping")
|
||||
continue
|
||||
|
||||
output_data = []
|
||||
total_batches = (len(encoded_texts) + batch_size - 1) // batch_size
|
||||
|
||||
for i in tqdm.tqdm(
|
||||
range(0, len(encoded_texts), batch_size),
|
||||
total=total_batches,
|
||||
desc="Computing perplexity",
|
||||
):
|
||||
batch_encoded = encoded_texts[i : i + batch_size]
|
||||
batch_texts = texts[i : i + batch_size]
|
||||
|
||||
# Find max length in batch and pad
|
||||
max_len = max(len(seq) for seq in batch_encoded)
|
||||
padded_ids = []
|
||||
masks = []
|
||||
|
||||
for seq in batch_encoded:
|
||||
pad_len = max_len - len(seq)
|
||||
padded_seq = seq + [tokenizer.pad_id] * pad_len
|
||||
mask = [True] * len(seq) + [False] * pad_len
|
||||
padded_ids.append(padded_seq)
|
||||
masks.append(mask)
|
||||
|
||||
# Convert to tensors
|
||||
input_ids = torch.tensor(padded_ids, device="cuda", dtype=torch.long)
|
||||
input_mask = torch.tensor(masks, device="cuda", dtype=torch.bool)
|
||||
|
||||
# Compute perplexity
|
||||
output = model(input_ids, input_mask=input_mask)
|
||||
logits = output["logits"]
|
||||
|
||||
# Shift for causal language modeling
|
||||
shifted_logits = logits[:, :-1, :] # [batch_size, seq_len-1, vocab_size]
|
||||
shifted_input_ids = input_ids[:, 1:] # [batch_size, seq_len-1]
|
||||
shifted_mask = input_mask[:, 1:] # [batch_size, seq_len-1]
|
||||
|
||||
# Compute cross entropy loss
|
||||
loss = F.cross_entropy(
|
||||
shifted_logits.flatten(0, 1),
|
||||
shifted_input_ids.flatten(0, 1),
|
||||
reduction="none",
|
||||
token_output = (
|
||||
os.path.join(output_dir, f"{label}_tokens.jsonl") if token_level else None
|
||||
)
|
||||
|
||||
loss = loss.view(shifted_input_ids.shape) # [batch_size, seq_len-1]
|
||||
loss = loss * shifted_mask
|
||||
sentence_loss = loss.sum(dim=1) / shifted_mask.sum(dim=1).clamp(min=1)
|
||||
perplexity = torch.exp(sentence_loss) # [batch_size]
|
||||
stats = process_file(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
items=items,
|
||||
text_key=text_key,
|
||||
batch_size=batch_size,
|
||||
max_length=max_length,
|
||||
token_level=token_level,
|
||||
max_samples=max_samples,
|
||||
output_file=token_output,
|
||||
label=label,
|
||||
)
|
||||
all_stats[label] = stats
|
||||
print_stats(label, stats)
|
||||
|
||||
for text, ppl in zip(batch_texts, perplexity):
|
||||
output_data.append({text_key: text, "ppl": float(ppl.item())})
|
||||
if token_output:
|
||||
print(f" token-level output: {token_output}")
|
||||
|
||||
# Write results
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
for item in output_data:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
|
||||
print(f"Perplexity computation complete. Results saved to {output_file}")
|
||||
summary_path = os.path.join(output_dir, "summary.json")
|
||||
with open(summary_path, "w", encoding="utf-8") as f:
|
||||
json.dump(all_stats, f, ensure_ascii=False, indent=2)
|
||||
print(f"\nSummary saved to {summary_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Perplexity evaluation on JSONL text.")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Perplexity and token-level loss evaluation on JSONL/JSON data."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--param_path", type=str, required=True, help="Path to the model directory."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_file", type=str, required=True, help="Path to the input file."
|
||||
"--input_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to input file, glob pattern, or directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_file", type=str, required=True, help="Path to the output file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size", type=int, default=4, help="Batch size for evaluation."
|
||||
"--output_dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Directory for output files (summary.json + per-file token JSONL).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text_key",
|
||||
|
|
@ -105,7 +409,37 @@ if __name__ == "__main__":
|
|||
default="text",
|
||||
help="Key for the text field in the input data.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size", type=int, default=4, help="Batch size for evaluation."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_length",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="Maximum sequence length (tokens). Longer sequences are truncated.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--token_level",
|
||||
action="store_true",
|
||||
help="Store per-token log_probs and token type analysis. "
|
||||
"Default: off (only aggregate stats).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_samples",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Maximum number of samples per file (random subsample). Default: all.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
with torch.inference_mode():
|
||||
process_file(**vars(args))
|
||||
main(
|
||||
param_path=args.param_path,
|
||||
input_path=args.input_path,
|
||||
output_dir=args.output_dir,
|
||||
text_key=args.text_key,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_length,
|
||||
token_level=args.token_level,
|
||||
max_samples=args.max_samples,
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue