# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. from abc import ABC from typing import List, Optional, Tuple, Union import torch from flow_matching.loss import MixturePathGeneralizedKL, EditFlowsLoss, EditFlowsLossReParam from flow_matching.path import MixtureDiscreteProbPath # for the scheduler only from flow_matching.path.scheduler import PolynomialConvexScheduler from flow_matching.path.editflows_adapter import EditFlowsPathAdapter # <-- NEW import from torch import Tensor from torch.nn.modules.loss import _Loss class SourceDistribution(ABC): def __init__(self) -> None: ... def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: ... def sample_like(self, tensor_like: Tensor) -> Tensor: ... class MaskedSourceDistribution(SourceDistribution): def __init__(self, mask_token: int) -> None: self.mask_token = mask_token @property def masked(self) -> bool: return True def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: return torch.zeros(tensor_size, device=device).fill_(self.mask_token).long() def sample_like(self, tensor_like: Tensor) -> Tensor: return torch.zeros_like(tensor_like).fill_(self.mask_token).long() import torch from typing import List, Tuple, Optional, Union Tensor = torch.Tensor class UniformSourceDistribution: def __init__(self, vocab_size, special_token_ids = None): self.vocab_size = vocab_size self.special_token_ids = set(special_token_ids) if special_token_ids is not None else set() # Compute allowed tokens by removing all special tokens from vocab self._allowed_tokens = [i for i in range(vocab_size) if i not in self.special_token_ids] if len(self._allowed_tokens) == 0: raise ValueError(f"All tokens are special tokens: {special_token_ids}") @property def masked(self) -> bool: return False @torch.no_grad() def _sample_from_allowed(self, shape, device, allowed_tokens = None, generator = None): """Sample uniformly from allowed tokens with given shape.""" if allowed_tokens is None: allowed_tokens = self._allowed_tokens if len(allowed_tokens) == 0: raise ValueError("No allowed tokens provided") allowed_tensor = torch.tensor(allowed_tokens, dtype=torch.long, device=device) num_allowed = len(allowed_tokens) if shape.numel() == 0: # Return an empty tensor (length-0 sequence) return torch.empty(shape, dtype=torch.long, device=device) # Sample indices in [0, num_allowed) indices = torch.randint( low=0, high=num_allowed, size=shape, device=device, generator=generator ) return allowed_tensor[indices] @torch.no_grad() def sample(self, tensor_size, device, allowed_tokens = None, generator = None): return self._sample_from_allowed(tensor_size, device, allowed_tokens, generator) @torch.no_grad() def sample_like(self, tensor_like, allowed_tokens = None, generator = None): """ Keep original semantics: sample tokens with the SAME SHAPE(S) as tensor_like. """ if isinstance(tensor_like, (list, tuple)): return [ self._sample_from_allowed(seq.shape, seq.device, allowed_tokens, generator) for seq in tensor_like ] return self._sample_from_allowed(tensor_like.shape, tensor_like.device, allowed_tokens, generator) @torch.no_grad() def sample_like( self, tensor_like, allowed_tokens: Optional[List[int]] = None, min_len_factor: float = 0.0, max_len_factor: float = 2.0, generator: Optional[torch.Generator] = None, ): """ For each reference x1, sample x0 with LENGTH L ~ Uniform{ floor(min_len_factor*N) .. floor(max_len_factor*N) }, defaulting to [0, 2*N]. Returns an empty tensor if L == 0. Supports Tensor (1D) or List[Tensor] (ragged). If you pass a 2D tensor, we interpret N as the last dimension and return a 1D sequence for that tensor. """ def _one(x1): assert x1.dim() >= 1, "x1 must be at least 1D" device = x1.device dtype = torch.long N = int(x1.shape[-1]) # use last-dim length as reference lo = int(max(0, int(min_len_factor * N))) hi = int(max(0, int(max_len_factor * N))) # randint is [low, high), so make high inclusive with +1 L = int(torch.randint(low=lo, high=hi + 1, size=(1,), device=device, generator=generator).item()) if L == 0: return torch.empty((0,), dtype=dtype, device=device) return self._sample_from_allowed(torch.Size([L]), device, allowed_tokens, generator) if isinstance(tensor_like, (list, tuple)): return [_one(seq) for seq in tensor_like] else: return _one(tensor_like) def sample_x0_from_x1(self, x1, pad_id, allowed_tokens, scale_size = 1.0, bos_id = 0, eos_id = 2): """ For each sequence in x1, sample an x0 whose *core* length (excluding BOS/EOS) is in [0, scale_size * len_valid(x1)], where len_valid(x1) counts only tokens that are NOT {BOS, EOS, PAD}. Rules: - x0 always starts with BOS and ends with EOS - x0 core tokens are sampled uniformly from vocab excluding {BOS, EOS, PAD} - We batch and pad x0 to a common length (B, L0) with pad_id - "length of x0 does not account for BOS and EOS" = the sampled core length Returns: x0: (B, L0) Long, padded with pad_id """ device = x1.device B, L1 = x1.shape # compute valid length of x1 per sequence # valid = not pad, not BOS, not EOS valid_mask_x1 = (x1 != pad_id) & (x1 != bos_id) & (x1 != eos_id) valid_len = valid_mask_x1.sum(dim=1) # (B,) # we will store all sequences here before padding x0_seqs = [] for b in range(B): max_core_len = int((1 + scale_size) * valid_len[b].item()) # may be 0 min_core_len = int((1 - scale_size) * valid_len[b].item()) # sample core length in [0, max_core_len] core_len = int(torch.randint(low=min_core_len, high=max_core_len + 1, size=(1,), device=device).item()) # 4) sample core tokens if core_len > 0: idx = torch.randint(0, allowed_tokens.size(0), (core_len,), device=device) core_tokens = allowed_tokens[idx] # (core_len,) else: core_tokens = torch.empty(0, dtype=torch.long, device=device) # 5) build full x0: [BOS] + core + [EOS] seq = torch.cat([ torch.tensor([bos_id], device=device, dtype=torch.long), core_tokens, torch.tensor([eos_id], device=device, dtype=torch.long), ], dim=0) # (1 + core_len + 1,) x0_seqs.append(seq) x0 = torch.nn.utils.rnn.pad_sequence(x0_seqs, batch_first=True, padding_value=pad_id) return x0 # class UniformSourceDistribution(SourceDistribution): # def __init__(self, vocab_size: int, special_token_ids: Optional[List[int]] = None) -> None: # self.vocab_size = vocab_size # self.special_token_ids = set(special_token_ids) if special_token_ids is not None else set() # # Compute allowed tokens by removing all special tokens from vocab # self._allowed_tokens = [i for i in range(vocab_size) if i not in self.special_token_ids] # if len(self._allowed_tokens) == 0: # raise ValueError(f"All tokens are special tokens: {special_token_ids}") # @property # def masked(self) -> bool: # return False # def _sample_from_allowed( # self, shape: torch.Size, device: torch.device, allowed_tokens: Optional[List[int]] = None # ) -> Tensor: # """Sample uniformly from allowed tokens.""" # if allowed_tokens is None: # # Use default allowed tokens (vocab minus special tokens) # allowed_tokens = self._allowed_tokens # if len(allowed_tokens) == 0: # raise ValueError(f"No allowed tokens provided") # # print(f"allowed_tokens: {allowed_tokens}") # # Sample indices into allowed_tokens, then map to actual token IDs # allowed_tensor = torch.tensor(allowed_tokens, dtype=torch.long, device=device) # num_allowed = len(allowed_tokens) # # Sample indices in [0, num_allowed) # indices = torch.randint(size=shape, high=num_allowed, device=device) # return allowed_tensor[indices] # def sample( # self, # tensor_size: Tuple[int, ...], # device: torch.device, # allowed_tokens: Optional[List[int]] = None # ) -> Tensor: # return self._sample_from_allowed(tensor_size, device, allowed_tokens) # def sample_like( # self, # tensor_like: Union[Tensor, List[Tensor]], # allowed_tokens: Optional[List[int]] = None # ) -> Union[Tensor, List[Tensor]]: # """ # Sample uniform tokens matching the shape of tensor_like. # Supports both Tensor and List[Tensor] for ragged inputs. # Args: # tensor_like: Either a Tensor or List[Tensor] (for ragged inputs) # allowed_tokens: Optional list of allowed token IDs. If None, uses vocab minus special tokens. # Returns: # Tensor or List[Tensor] matching the input shape(s) # """ # # Handle ragged input (list of tensors) # if isinstance(tensor_like, (list, tuple)): # return [ # self._sample_from_allowed(seq.shape, seq.device, allowed_tokens) # for seq in tensor_like # ] # # Handle regular tensor input # return self._sample_from_allowed(tensor_like.shape, tensor_like.device, allowed_tokens) class EmpiricalSourceDistribution(SourceDistribution): def __init__(self, vocab_size: int, probs: torch.Tensor, length: int): self.vocab_size = vocab_size self.registered = probs / probs.sum() self.length = length # e.g., 100 tokens @property def masked(self) -> bool: return False def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: B = tensor_size[0] idx = torch.multinomial(self.registered.to(device), num_samples=self.length, replacement=True) return idx.view(1, self.length).repeat(B, 1) def sample_like(self, tensor_like: Tensor) -> Tensor: B = tensor_like.shape[0] idx = torch.multinomial(self.registered.to(tensor_like.device), num_samples=self.length, replacement=True) return idx.view(1, self.length).repeat(B, 1) # NOTE: return type changed to the adapter (ragged). We only rely on .scheduler(t) and .sample(...). def get_path(scheduler_type: str, exponent: Optional[float] = None, eps_id: int = -1) -> EditFlowsPathAdapter: if scheduler_type == "polynomial": # paper uses cubic => exponent=3 scheduler = PolynomialConvexScheduler(n=exponent) else: raise ValueError(f"{scheduler_type} is not supported") # MixtureDiscreteProbPath carries the scheduler; the adapter will sample ragged z_t itself mixture = MixtureDiscreteProbPath(scheduler=scheduler) return EditFlowsPathAdapter(mixture_path=mixture, eps_id=eps_id) def get_source_distribution( source_distribution: str, p_emp: Optional[Tensor] = None, length: Optional[int] = None, vocab_size: Optional[int] = None, special_token_ids: Optional[List[int]] = None, ) -> SourceDistribution: if p_emp is not None: assert vocab_size is not None and length is not None, "Empirical source requires vocab_size and length" return EmpiricalSourceDistribution(vocab_size=vocab_size, probs=p_emp, length=length) if source_distribution == "mask": assert vocab_size is not None, "Masked source requires vocab_size" return MaskedSourceDistribution(mask_token=vocab_size) elif source_distribution == "uniform": assert vocab_size is not None, "Uniform source requires vocab_size" return UniformSourceDistribution(vocab_size=vocab_size, special_token_ids=special_token_ids) else: raise ValueError(f"{source_distribution} is not supported") def get_loss_function(loss_function: str, path: Optional[Union[MixtureDiscreteProbPath, EditFlowsPathAdapter]] = None) -> _Loss: if loss_function == "cross_entropy": return torch.nn.CrossEntropyLoss() elif loss_function == "generalized_kl": assert path is not None # Generalized KL still expects a (dense) path; fine for DFM experiments return MixturePathGeneralizedKL(path=path) elif loss_function == "editflows": # Ragged EF loss does NOT need the path; training.step precomputes the weight return EditFlowsLoss(reduction="mean") elif loss_function == "editflows_reparam": return EditFlowsLossReParam(reduction="mean") else: raise ValueError(f"{loss_function} is not supported") # # Copyright (c) Meta Platforms, Inc. and affiliates. # # 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. # from abc import ABC # from typing import Optional, Tuple # import torch # from flow_matching.loss import MixturePathGeneralizedKL, EditFlowsLoss # from flow_matching.path import MixtureDiscreteProbPath, ProbPath, EditFlowsPathAdapter # from flow_matching.path.scheduler import PolynomialConvexScheduler # from torch import Tensor # from torch.nn.modules.loss import _Loss # class SourceDistribution(ABC): # def __init__( # self, # ) -> None: # ... # def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: # ... # def sample_like(self, tensor_like: Tensor) -> Tensor: # ... # class MaskedSourceDistribution(SourceDistribution): # def __init__(self, mask_token: int) -> None: # self.mask_token = mask_token # @property # def masked(self) -> bool: # return True # def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: # return torch.zeros(tensor_size, device=device).fill_(self.mask_token).long() # def sample_like(self, tensor_like: Tensor) -> Tensor: # return torch.zeros_like(tensor_like).fill_(self.mask_token).long() # class UniformSourceDistribution(SourceDistribution): # def __init__(self, vocab_size: int) -> None: # self.vocab_size = vocab_size # @property # def masked(self) -> bool: # return False # def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: # return torch.randint(size=tensor_size, high=self.vocab_size, device=device) # def sample_like(self, tensor_like: Tensor) -> Tensor: # return torch.randint_like(tensor_like, high=self.vocab_size) # class EmpiricalSourceDistribution(SourceDistribution): # def __init__(self, vocab_size: int, probs: torch.Tensor, length: int): # self.vocab_size = vocab_size # self.registered = probs / probs.sum() # self.length = length # e.g., 100 tokens as in the paper’s variant # @property # def masked(self) -> bool: return False # def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: # B = tensor_size[0] # idx = torch.multinomial(self.registered.to(device), num_samples=self.length, replacement=True) # return idx.view(1, self.length).repeat(B, 1) # def sample_like(self, tensor_like: Tensor) -> Tensor: # B = tensor_like.shape[0] # idx = torch.multinomial(self.registered.to(tensor_like.device), num_samples=self.length, replacement=True) # return idx.view(1, self.length).repeat(B, 1) # def get_path(scheduler_type: str, exponent: Optional[float] = None) -> ProbPath: # if scheduler_type == "polynomial": # scheduler = PolynomialConvexScheduler(n=exponent) # else: # raise ValueError(f"{scheduler_type} is not supported") # return EditFlowsPathAdapter(path=MixtureDiscreteProbPath(scheduler=scheduler)) # still need to decide (1) whether to implement the z_0, z_1 creation here, and (2) how to pass the eps_id to the adapter # def get_source_distribution( # source_distribution: str, p_emp: Optional[Tensor] = None, length: Optional[int] = None, vocab_size: int = None # ) -> SourceDistribution: # if p_emp is not None: # return EmpiricalSourceDistribution(vocab_size=vocab_size, probs=p_emp, length=length) # if source_distribution == "mask": # return MaskedSourceDistribution(mask_token=vocab_size) # elif source_distribution == "uniform": # return UniformSourceDistribution(vocab_size=vocab_size) # else: # raise ValueError(f"{source_distribution} is not supported") # def get_loss_function(loss_function: str, path: Optional[ProbPath] = None) -> _Loss: # if loss_function == "cross_entropy": # return torch.nn.CrossEntropyLoss() # elif loss_function == "generalized_kl": # assert path is not None # return MixturePathGeneralizedKL(path=path) # elif loss_function == "editflows": # assert path is not None # return EditFlowsLoss(path=path) # else: # raise ValueError(f"{loss_function} is not supported")