File size: 5,674 Bytes
7f316fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | # Copyright (c) Meta Platforms, Inc.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.
# Part of this implementation is adapted from https://github.com/louaaron/Score-Entropy-Discrete-Diffusion
# which is released under MIT license
import math
from collections import Counter
from typing import Iterable, List, Optional, Sequence, Tuple
import torch
import torch.nn.functional as F
from torch import Tensor
from torch.utils.data import DataLoader
from tqdm import tqdm
# ---------------------------
# Helpers for ragged handling
# ---------------------------
def _batch_iter(seq: Sequence, batch_size: int) -> Iterable[Sequence]:
for i in range(0, len(seq), batch_size):
yield seq[i : i + batch_size]
def _pad_ragged_to_dense(
seqs: Sequence[Tensor],
pad_id: int,
) -> Tuple[Tensor, Tensor]:
"""
Pad a list of 1-D LongTensors -> (B, Lmax) LongTensor and an attention mask (B, Lmax) bool.
pad_id is written into padded positions; attention_mask == 1 for real tokens.
"""
if len(seqs) == 0:
return torch.empty(0, 0, dtype=torch.long), torch.empty(0, 0, dtype=torch.bool)
device = seqs[0].device
lens = [int(s.numel()) for s in seqs]
Lmax = max(lens)
B = len(seqs)
out = torch.full((B, Lmax), pad_id, dtype=torch.long, device=device)
attn = torch.zeros((B, Lmax), dtype=torch.bool, device=device)
for i, s in enumerate(seqs):
L = s.numel()
if L > 0:
out[i, :L] = s
attn[i, :L] = True
return out, attn
# ---------------------------
# Perplexity (ragged-friendly)
# ---------------------------
@torch.no_grad()
def compute_perplexity(
samples: List[Tensor],
*,
lm_model, # a causal LM compatible with your tokenizer/vocab (e.g., HF AutoModelForCausalLM)
pad_id: int,
batch_size: int = 16,
) -> Tensor:
"""
Compute perplexity on ragged samples using a provided *compatible* causal LM.
Args:
samples: list of 1-D LongTensors (ragged, token ids in the *same* vocab the LM expects)
lm_model: a causal LM (e.g., transformers.AutoModelForCausalLM) on the right device
pad_id: integer id used for padding during evaluation
batch_size: eval batch size
Returns:
Scalar tensor: mean perplexity across sequences.
"""
if len(samples) == 0:
return torch.tensor(float("nan"))
device = samples[0].device
lm_model.eval()
ppl_vals: List[Tensor] = []
for chunk in _batch_iter(samples, batch_size):
# pad ragged batch
x, attn = _pad_ragged_to_dense(chunk, pad_id=pad_id) # (B, L), (B, L)
# shift for next-token prediction
x_in = x[:, :-1]
x_tgt = x[:, 1:]
attn_in = attn[:, :-1] # (B, L-1)
# forward LM
out = lm_model(input_ids=x_in, attention_mask=attn_in, use_cache=False)
logits = out.logits # (B, L-1, V)
# compute token-wise NLL with ignore on pads
# set ignored targets to -100 per HF convention
labels = x_tgt.clone()
labels[~attn[:, 1:]] = -100
# cross-entropy averaged over *valid* tokens; we want per-sequence mean → then exp
ce = F.cross_entropy(
logits.transpose(-1, -2), # (B, V, L-1)
labels,
reduction="none",
ignore_index=-100,
) # (B, L-1)
# per-sequence mean over valid positions
valid = (labels != -100).float()
token_counts = valid.sum(dim=1).clamp_min(1.0) # avoid div-by-zero
ce_seq = (ce * valid).sum(dim=1) / token_counts # (B,)
ppl_seq = ce_seq.exp() # (B,)
ppl_vals.append(ppl_seq)
ppl_all = torch.cat(ppl_vals, dim=0) # (N,)
return ppl_all.mean()
# ---------------------------
# Entropy (ragged-friendly)
# ---------------------------
def _sample_entropy_1d(sample: Tensor) -> float:
"""
Entropy (base-2) of a 1-D LongTensor by empirical token histogram.
"""
if sample.numel() == 0:
return 0.0
histogram = Counter(sample.tolist())
total = float(sum(histogram.values()))
ent = 0.0
for c in histogram.values():
p = c / total
ent -= p * math.log2(max(p, 1e-12))
return ent
@torch.no_grad()
def compute_entropy(samples: List[Tensor]) -> Tensor:
"""
Mean per-sequence token entropy (base-2) for ragged samples.
"""
if len(samples) == 0:
return torch.tensor(float("nan"))
device = samples[0].device
vals = [_sample_entropy_1d(s) for s in samples]
return torch.tensor(sum(vals) / len(vals), device=device)
# ---------------------------
# Likelihood / ELBO (EF note)
# ---------------------------
@torch.no_grad()
def estimate_likelihood(
model: torch.nn.Module,
dataloader: DataLoader,
source_distribution,
path,
n_discretization: int,
device: torch.device,
batch_size: int = 32,
epsilon: float = 1e-3,
) -> Tensor:
"""
Placeholder: ELBO-style likelihood estimation used for DFM does not
directly apply to Edit Flows (different objective / dynamics).
If you need a quantitative likelihood-like diagnostic for EF,
consider CTMC pathwise estimators or reverse-time simulators instead.
We raise NotImplementedError to avoid silently reporting a mismatched metric.
"""
raise NotImplementedError(
"ELBO/likelihood estimation used for discrete flow matching is not applicable to Edit Flows. "
"Use task metrics or CTMC-based diagnostics instead."
)
|