| import math |
| import statistics |
| from typing import Dict, List, Union |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| import sacrebleu |
| import transformers |
| from tqdm import tqdm |
|
|
| from utils.logging_utils import log_for_0 |
|
|
|
|
| |
| |
| |
| def _mean_std_sem(values): |
| n = len(values) |
| mean = sum(values) / n |
| std = statistics.pstdev(values) if n > 1 else 0.0 |
| sem = std / math.sqrt(n) if n > 1 else 0.0 |
| return mean, std, sem |
|
|
|
|
| def compute_bleu(hypotheses, references): |
| return sacrebleu.corpus_bleu(hypotheses, [references], lowercase=True, use_effective_order=True).score |
|
|
|
|
| def compute_rouge(hypotheses, references, return_std=False): |
| from rouge_score import rouge_scorer |
| scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True) |
| r1, r2, rL = [], [], [] |
| for hyp, ref in zip(hypotheses, references): |
| s = scorer.score(ref, hyp) |
| r1.append(s["rouge1"].fmeasure * 100) |
| r2.append(s["rouge2"].fmeasure * 100) |
| rL.append(s["rougeL"].fmeasure * 100) |
| m1, s1, e1 = _mean_std_sem(r1) |
| m2, s2, e2 = _mean_std_sem(r2) |
| mL, sL, eL = _mean_std_sem(rL) |
| means = {"rouge1": m1, "rouge2": m2, "rougeL": mL} |
| if not return_std: |
| return means |
| stds = { |
| "rouge1_std": s1, "rouge2_std": s2, "rougeL_std": sL, |
| "rouge1_sem": e1, "rouge2_sem": e2, "rougeL_sem": eL, |
| } |
| return means, stds |
|
|
|
|
| |
| |
| |
| class NLL: |
| """PyTorch implementation of NLL metric.""" |
|
|
| def __init__(self): |
| self.reset() |
|
|
| def reset(self): |
| self.mean_value = torch.tensor(0.0, dtype=torch.float64) |
| self.weight = torch.tensor(0.0, dtype=torch.float64) |
|
|
| def update(self, value: Union[float, torch.Tensor], weight: Union[float, torch.Tensor] = 1.0): |
| if not isinstance(value, torch.Tensor): |
| value = torch.tensor(value, dtype=torch.float64) |
| if not isinstance(weight, torch.Tensor): |
| weight = torch.tensor(weight, dtype=torch.float64) |
| weight = torch.broadcast_to(weight, value.shape) |
| if value.numel() == 0: |
| return |
| self.mean_value = self.mean_value + value.sum() |
| self.weight = self.weight + weight.sum() |
|
|
|
|
| class Perplexity(NLL): |
| def compute(self) -> torch.Tensor: |
| return torch.exp(self.mean_value / self.weight) |
|
|
|
|
| class MeanMetric: |
| def __init__(self): |
| self.reset() |
|
|
| def reset(self): |
| self.sum_value = torch.tensor(0.0, dtype=torch.float64) |
| self.count = torch.tensor(0.0, dtype=torch.float64) |
|
|
| def update(self, value: Union[float, torch.Tensor]): |
| if not isinstance(value, torch.Tensor): |
| value = torch.tensor(value, dtype=torch.float64) |
| self.sum_value = self.sum_value + value.sum() |
| self.count = self.count + value.numel() |
|
|
| def compute(self) -> torch.Tensor: |
| return self.sum_value / self.count |
|
|
|
|
| class Metrics: |
| def __init__( |
| self, |
| gen_ppl_eval_model_name_or_path=None, |
| eval_ppl_batch_size=None, |
| eval_context_size=1024, |
| ) -> None: |
| self.gen_ppl = Perplexity() |
| self.sample_entropy = MeanMetric() |
| self.eval_ppl_batch_size = eval_ppl_batch_size |
| self.gen_ppl_eval_model_name_or_path = gen_ppl_eval_model_name_or_path |
| self.eval_context_size = eval_context_size |
| self._eval_model = None |
| self._eval_device = None |
|
|
| |
| use_fast = "mt5" not in gen_ppl_eval_model_name_or_path.lower() |
| self.tokenizer = transformers.AutoTokenizer.from_pretrained( |
| gen_ppl_eval_model_name_or_path, use_fast=use_fast, |
| ) |
| if self.tokenizer.pad_token is None: |
| self.tokenizer.pad_token = self.tokenizer.eos_token |
| self.tokenizer.pad_token_id = self.tokenizer.eos_token_id |
|
|
| def reset(self): |
| self.gen_ppl.reset() |
| self.sample_entropy.reset() |
|
|
| def _eval_retokenize(self, text_samples, max_length): |
| out = self.tokenizer( |
| text_samples, |
| return_tensors="np", |
| return_token_type_ids=False, |
| return_attention_mask=True, |
| truncation=True, |
| padding=True, |
| max_length=max_length, |
| ) |
| return out["input_ids"], out["attention_mask"], self.eval_context_size |
|
|
| @torch.no_grad() |
| def _compute_batch_nlls(self, input_ids: torch.Tensor, attention_mask: torch.Tensor): |
| model = self._eval_model |
| eos_token_id = self.tokenizer.eos_token_id |
| logits = model(input_ids=input_ids, attention_mask=attention_mask).logits |
| targets = input_ids[:, 1:] |
| logits_pred = logits[:, :-1, :] |
| log_normalizers = torch.logsumexp(logits_pred.to(torch.float32), dim=-1) |
| target_logits = logits_pred.gather(-1, targets.unsqueeze(-1)).squeeze(-1).to(torch.float32) |
| nlls = log_normalizers - target_logits |
| is_eos = (input_ids == eos_token_id) |
| first_eos = (is_eos.to(torch.int32).cumsum(dim=-1) == 1) |
| token_mask = (input_ids != eos_token_id) |
| valid_tokens = first_eos[:, 1:].to(torch.int32) + token_mask[:, 1:].to(torch.int32) |
| return nlls, valid_tokens |
|
|
| def record_generative_perplexity( |
| self, |
| text_samples: List[str], |
| max_length: int, |
| retokenize: bool = True, |
| ) -> Dict: |
| import os |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" |
|
|
| if self._eval_model is None: |
| from transformers import AutoModelForCausalLM |
| log_for_0(f"Loading PyTorch model: {self.gen_ppl_eval_model_name_or_path}") |
| self._eval_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self._eval_model = AutoModelForCausalLM.from_pretrained( |
| self.gen_ppl_eval_model_name_or_path, |
| torch_dtype=torch.bfloat16, |
| ).to(self._eval_device).eval() |
| log_for_0("PPL model cached for reuse") |
|
|
| device = self._eval_device |
|
|
| if retokenize: |
| samples, attn_mask, eval_context_size = self._eval_retokenize(text_samples, max_length=max_length) |
| else: |
| samples = text_samples |
| attn_mask = np.ones(samples.shape) |
| eval_context_size = samples.shape[-1] |
|
|
| batch_size = self.eval_ppl_batch_size or samples.shape[0] |
| batch_size = min(batch_size, samples.shape[0]) or 1 |
| num_batches = (samples.shape[0] + batch_size - 1) // batch_size |
| log_for_0(f"PPL: batch_size={batch_size}, {num_batches} batches") |
|
|
| per_sample_nll_sum = np.zeros(samples.shape[0], dtype=np.float64) |
| per_sample_token_count = np.zeros(samples.shape[0], dtype=np.float64) |
|
|
| for i in tqdm(range(num_batches), desc="Evaluating perplexity"): |
| batch_start = i * batch_size |
| batch_end = min((i + 1) * batch_size, samples.shape[0]) |
| batch_samples = samples[batch_start:batch_end] |
| batch_attn_mask = attn_mask[batch_start:batch_end] |
|
|
| for chunk_start in range(0, batch_samples.shape[1], eval_context_size): |
| chunk_end = min(chunk_start + eval_context_size, batch_samples.shape[1]) |
| sample_chunk = batch_samples[:, chunk_start:chunk_end] |
| attn_mask_chunk = batch_attn_mask[:, chunk_start:chunk_end] |
|
|
| input_ids = torch.from_numpy(sample_chunk).to(device).long() |
| attn = torch.from_numpy(attn_mask_chunk).to(device).long() |
| nlls, valid_tokens = self._compute_batch_nlls(input_ids, attn) |
|
|
| nlls_np = nlls.detach().cpu().numpy().astype(np.float64) |
| valid_tokens_np = valid_tokens.detach().cpu().numpy().astype(np.float64) |
| weighted_nlls = nlls_np * valid_tokens_np |
|
|
| self.gen_ppl.update(torch.from_numpy(weighted_nlls), |
| torch.from_numpy(valid_tokens_np)) |
|
|
| per_sample_nll_sum[batch_start:batch_end] += weighted_nlls.sum(axis=-1) |
| per_sample_token_count[batch_start:batch_end] += valid_tokens_np.sum(axis=-1) |
|
|
| with np.errstate(divide="ignore", invalid="ignore"): |
| per_sample_ppl = np.exp(per_sample_nll_sum / per_sample_token_count) |
| per_sample_ppl = np.where(per_sample_token_count > 0, per_sample_ppl, np.nan).tolist() |
|
|
| per_sample_entropy = [] |
| for i in range(samples.shape[0]): |
| valid_len = int(attn_mask[i].sum()) |
| valid_tokens = samples[i, :valid_len] |
| _, counts = np.unique(valid_tokens, return_counts=True) |
| probs = counts.astype(np.float32) / counts.sum() |
| entropy = float(-np.sum(probs * np.log(probs + 1e-10))) |
| per_sample_entropy.append(entropy) |
| self.sample_entropy.update(entropy) |
|
|
| return { |
| "ppl": float(self.gen_ppl.compute()), |
| "per_sample_ppl": per_sample_ppl, |
| "mean_entropy": sum(per_sample_entropy) / len(per_sample_entropy), |
| } |
|
|
