446 lines
15 KiB
Python
446 lines
15 KiB
Python
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
|
|
import tqdm
|
|
|
|
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(
|
|
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],
|
|
):
|
|
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()
|
|
|
|
input_files = _collect_input_files(input_path)
|
|
if not input_files:
|
|
print(f"No input files found at {input_path}")
|
|
return
|
|
|
|
print(f"Found {len(input_files)} file(s) to evaluate")
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
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
|
|
|
|
token_output = (
|
|
os.path.join(output_dir, f"{label}_tokens.jsonl") if token_level else None
|
|
)
|
|
|
|
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)
|
|
|
|
if token_output:
|
|
print(f" token-level output: {token_output}")
|
|
|
|
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 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_path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to input file, glob pattern, or directory.",
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir",
|
|
type=str,
|
|
required=True,
|
|
help="Directory for output files (summary.json + per-file token JSONL).",
|
|
)
|
|
parser.add_argument(
|
|
"--text_key",
|
|
type=str,
|
|
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():
|
|
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,
|
|
)
|