Spaces:
Running
Running
| # coding: utf-8 | |
| import torch | |
| import torch.nn.functional as F | |
| from typing import Any, Tuple, List | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| from constants import TARGET_PAD | |
| class Batch: | |
| """ | |
| Batch wrapper that normalizes different batch formats. | |
| Accepts either: | |
| - legacy torchtext batch object with attributes: | |
| .src (tensor), .src_lengths (tensor or int), .trg (tensor), .file_paths (list) | |
| - modern tuple from DataLoader/collate_fn: | |
| (src_padded, src_lengths, trg_padded, files) | |
| Args: | |
| torch_batch: batch object or tuple from DataLoader | |
| pad_index: integer index used for source padding | |
| model: model instance (used to check use_cuda) | |
| """ | |
| def __init__(self, torch_batch: Any, pad_index: int, model: Any): | |
| # Initialize common attributes | |
| self.src = None # LongTensor [B, S] | |
| self.src_lengths = None # LongTensor [B] | |
| self.src_mask = None # Bool/ByteTensor [B,1,S] | |
| self.nseqs = 0 | |
| self.trg = None # FloatTensor [B, T, trg_size] | |
| self.trg_input = None # FloatTensor [B, T, trg_size] (for teacher forcing / model input) | |
| self.trg_mask = None # Bool/ByteTensor [B,1,T] (True where not padded) | |
| self.trg_lengths = None # int or LongTensor | |
| self.ntokens = 0 # number of non-pad target frames (sum) | |
| self.file_paths: List[str] = [] | |
| # model flags | |
| self.use_cuda = getattr(model, "use_cuda", False) | |
| self.target_pad = TARGET_PAD | |
| # Unpack depending on batch format | |
| # Handle legacy torchtext-like object with attributes | |
| if hasattr(torch_batch, "src") and hasattr(torch_batch, "file_paths"): | |
| # torchtext-style Example batch | |
| try: | |
| # src may be (src_tensor, src_lengths) or src_tensor alone | |
| if isinstance(torch_batch.src, tuple) or isinstance(torch_batch.src, list): | |
| self.src, self.src_lengths = torch_batch.src | |
| else: | |
| self.src = torch_batch.src | |
| # try to obtain lengths if available on batch | |
| self.src_lengths = getattr(torch_batch, "src_lengths", torch.tensor([s.size(0) for s in self.src], dtype=torch.long)) | |
| except Exception: | |
| # fallback: assume src is tensor and compute lengths from padding | |
| self.src = torch_batch.src | |
| self.src_lengths = getattr(torch_batch, "src_lengths", torch.sum(self.src != pad_index, dim=1)) | |
| self.file_paths = list(getattr(torch_batch, "file_paths", [])) | |
| # Targets (if present) | |
| if hasattr(torch_batch, "trg"): | |
| self.trg = torch_batch.trg | |
| # Handle tuple produced by DataLoader / collate_fn | |
| elif isinstance(torch_batch, (tuple, list)) and len(torch_batch) >= 3: | |
| # Expected format: (src_padded, src_lengths, trg_padded, files) | |
| # Some collate_fns may not return src_lengths; handle both cases. | |
| # Common expected: | |
| # src_padded: LongTensor [B, S] | |
| # src_lengths: LongTensor [B] | |
| # trg_padded: FloatTensor [B, T, trg_size] | |
| # files: list[str] | |
| try: | |
| self.src = torch_batch[0] | |
| self.src_lengths = torch_batch[1] | |
| self.trg = torch_batch[2] | |
| # files may be absent or None | |
| if len(torch_batch) > 3: | |
| self.file_paths = list(torch_batch[3]) | |
| else: | |
| self.file_paths = [] | |
| except Exception: | |
| raise ValueError("Unrecognized tuple batch format. Expected (src, src_lengths, trg, files).") | |
| else: | |
| raise ValueError("Unrecognized batch format passed to Batch.") | |
| # Ensure shapes / dtypes | |
| if isinstance(self.src_lengths, int): | |
| self.src_lengths = torch.tensor([self.src_lengths] * self.src.size(0), dtype=torch.long) | |
| if self.src is not None and not isinstance(self.src_lengths, torch.Tensor): | |
| # attempt to compute lengths from padding if possible | |
| try: | |
| self.src_lengths = torch.sum(self.src != pad_index, dim=1).to(torch.long) | |
| except Exception: | |
| # fallback zeros | |
| self.src_lengths = torch.zeros(self.src.size(0), dtype=torch.long) | |
| # src_mask: True where token != pad_index | |
| if self.src is not None: | |
| self.src_mask = (self.src != pad_index).unsqueeze(1) # [B,1,S] | |
| self.nseqs = self.src.size(0) | |
| # Targets handling | |
| if self.trg is not None: | |
| # trg is expected shape [B, T, trg_size] | |
| # trg_lengths: number of frames (T) - if not available infer from shape | |
| try: | |
| self.trg_lengths = self.trg.shape[1] | |
| except Exception: | |
| self.trg_lengths = None | |
| # trg_input: the model expects target input. Keep same shape as trg. | |
| # If you want to shift / remove last frame, do that in the training loop or here by uncommenting the next line: | |
| # self.trg_input = self.trg[:, :-1, :].clone() | |
| self.trg_input = self.trg.clone() | |
| # trg mask: True where frame is not padding. We assume padding frames have all elements equal to TARGET_PAD. | |
| # To detect padded frames, compare first element of each frame to TARGET_PAD (fast and typical). | |
| # Fallback: compare sum across frame (if consistent). | |
| try: | |
| # If last dim exists | |
| if self.trg.dim() == 3: | |
| self.trg_mask = (self.trg[:, :, 0] != self.target_pad).unsqueeze(1) # [B,1,T] | |
| else: | |
| # unexpected dimensions -> assume all valid | |
| self.trg_mask = torch.ones((self.trg.size(0), 1, self.trg.size(1)), dtype=torch.bool) | |
| except Exception: | |
| # fallback: assume all frames valid | |
| self.trg_mask = torch.ones((self.trg.size(0), 1, self.trg.size(1)), dtype=torch.bool) | |
| # ntokens: number of non-padding frames across the batch | |
| try: | |
| self.ntokens = int(torch.sum(self.trg_mask).item()) | |
| except Exception: | |
| self.ntokens = 0 | |
| # Filepaths: ensure list | |
| if self.file_paths is None: | |
| self.file_paths = [] | |
| # Move to GPU if required | |
| if self.use_cuda: | |
| self._make_cuda() | |
| def _make_cuda(self): | |
| """Move the batch tensors to GPU (in-place).""" | |
| if self.src is not None: | |
| self.src = self.src.to(device) | |
| self.src_mask = self.src_mask.to(device) | |
| self.src_lengths = self.src_lengths.to(device) | |
| if self.trg_input is not None: | |
| self.trg_input = self.trg_input.to(device) | |
| if self.trg is not None: | |
| self.trg = self.trg.to(device) | |
| if self.trg_mask is not None: | |
| self.trg_mask = self.trg_mask.to(device) | |
| def to(self, device_: torch.device): | |
| """Move batch to specified device and return self (convenience).""" | |
| if self.src is not None: | |
| self.src = self.src.to(device_) | |
| self.src_mask = self.src_mask.to(device_) | |
| self.src_lengths = self.src_lengths.to(device_) | |
| if self.trg_input is not None: | |
| self.trg_input = self.trg_input.to(device_) | |
| if self.trg is not None: | |
| self.trg = self.trg.to(device_) | |
| if self.trg_mask is not None: | |
| self.trg_mask = self.trg_mask.to(device_) | |
| return self | |