roughly equivalent to for item in loader yield inferitem params arguments loader torch utils data dataloader or any iterator the iterator that will be used to apply infer on infer any function the function to apply of each element of loader params dict the parameters passed to infer along with every item loaderbatchsize int optional if specified the items of loader are supposed to come as batch and are loaderbatched here making it roughly behave as for items in loader for i in loaderbatchsize item itemsi yield inferitem params self loader loader self infer infer self params params if loaderbatchsize 1 let s spare some time by deactivating altogether loaderbatchsize none self loaderbatchsize loaderbatchsize internal bookkeeping self loaderbatchindex none self loaderbatchdata none def lenself return lenself loader def iterself self iterator iterself loader return self def loaderbatchitemself if isinstanceself loaderbatchdata torch tensor batch data is simple tensor just fetch the slice result self loaderbatchdataself loaderbatchindex else batch data is assumed to be basemodeloutput or dict loaderbatched for k element in self loaderbatchdata items if isinstanceelement modeloutput convert modeloutput to tuple first element element totuple if isinstanceelement0 torch tensor loaderbatchedk tupleelself loaderbatchindex unsqueeze0 for el in element elif isinstanceelement0 np ndarray loaderbatchedk tuplenp expanddimselself loaderbatchindex 0 for el in element continue if k in hiddenstates pastkeyvalues attentions and isinstanceelement tuple those are stored as lists of tensors so need specific unbatching if isinstanceelement0 torch tensor loaderbatchedk tupleelself loaderbatchindex unsqueeze0 for el in element elif isinstanceelement0 np ndarray loaderbatchedk tuplenp expanddimselself loaderbatchindex 0 for el in element continue if element is none this can happen for optional data that get passed around loaderbatchedk none elif isinstanceelementself loaderbatchindex torch tensor take correct batch data but make it looked like batchsize1 for compatibility with other methods within transformers loaderbatchedk elementself loaderbatchindex unsqueeze0 elif isinstanceelementself loaderbatchindex np ndarray take correct batch data but make it looked like batchsize1 for compatibility with other methods within transformers loaderbatchedk np expanddimselementself loaderbatchindex 0 else this is typically a list so no need to unsqueeze loaderbatchedk elementself loaderbatchindex recreate the element by reusing the original class to make it look batchsize1 result self loaderbatchdata classloaderbatched self loaderbatchindex 1 return result def nextself if self loaderbatchindex is not none and self loaderbatchindex self loaderbatchsize we are currently unrolling a batch so we just need to return the current item within a batch return self loaderbatchitem we re out of items within a batch item nextself iterator processed self inferitem self params we now have a batch of inferred things if self loaderbatchsize is not none try to infer the size of the batch if isinstanceprocessed torch tensor firsttensor processed else key listprocessed keys0 firsttensor processedkey if isinstancefirsttensor list observedbatchsize lenfirsttensor else observedbatchsize firsttensor shape0 if 0 observedbatchsize self loaderbatchsize could be last batch so we can t unroll as many elements self loaderbatchsize observedbatchsize setting internal index to unwrap the batch self loaderbatchdata processed self loaderbatchindex 0 return self loaderbatchitem else we re not unrolling batches return processed class pipelinechunkiteratorpipelineiterator def initself loader infer params loaderbatchsizenone super initloader infer params def iterself self iterator iterself loader self subiterator none return self def nextself if self subiterator is none subiterator none means we haven t started a preprocess iterator so start it self subiterator self infernextself iterator self params try try to return next item processed nextself subiterator except stopiteration when a preprocess iterator ends we can start lookig at the next item chunkiterator will keep feeding until all elements of iterator all have created their subiterator and have been iterating against another way to look at it is we re basically flattening lists of lists into a single list but with generators self subiterator self infernextself iterator self params processed nextself subiterator return processed class pipelinepackiteratorpipelineiterator item is batched extremely similar to pipelineiterator in its unpacking mechanism but we have an extra required item which is the presence of islast that is because everything is flattened by pipelinechunkiterator we need to keep track of how to regroup here in the original process boundaries so that process and postprocess see the same data this iterator accumulates items possibly while unbatching until it its a islast and then just passes it on to the caller could be last batch so we can t unroll as many elements roughly equivalent to for item in loader yield infer item params arguments loader torch utils data dataloader or any iterator the iterator that will be used to apply infer on infer any function the function to apply of each element of loader params dict the parameters passed to infer along with every item loader_batch_size int optional if specified the items of loader are supposed to come as batch and are loader_batched here making it roughly behave as for items in loader for i in loader_batch_size item items i yield infer item params let s spare some time by deactivating altogether internal bookkeeping return item located at loader_batch_index within the current loader_batch_data batch data is simple tensor just fetch the slice batch data is assumed to be basemodeloutput or dict convert modeloutput to tuple first those are stored as lists of tensors so need specific unbatching this can happen for optional data that get passed around take correct batch data but make it looked like batch_size 1 for compatibility with other methods within transformers take correct batch data but make it looked like batch_size 1 for compatibility with other methods within transformers this is typically a list so no need to unsqueeze recreate the element by reusing the original class to make it look batch_size 1 we are currently unrolling a batch so we just need to return the current item within a batch we re out of items within a batch we now have a batch of inferred things try to infer the size of the batch could be last batch so we can t unroll as many elements setting internal index to unwrap the batch we re not unrolling batches roughly equivalent to for iterator in loader for item in iterator yield infer item params arguments loader torch utils data dataloader or any iterator the iterator that will be used to apply infer on infer any function the function to apply of each element of loader params dict the parameters passed to infer along with every item try to return next item when a preprocess iterator ends we can start lookig at the next item chunkiterator will keep feeding until all elements of iterator all have created their subiterator and have been iterating against another way to look at it is we re basically flattening lists of lists into a single list but with generators roughly equivalent to packed for item in loader packed append item if item is_last yield packed packed but it also handles cases where item are batched meaning it s a dict of tensor with first dimension 1 in that case it does packed for batch in loader item is batched for item in batch packed append item if item is_last yield packed packed arguments loader torch utils data dataloader or any iterator the iterator that will be used to apply infer on infer any function the function to apply of each element of loader params dict the parameters passed to infer along with every item loader_batch_size int optional if specified the items of loader are supposed to come as batch and are loader_batched here making it roughly behave as for items in loader for i in loader_batch_size item items i yield infer item params extremely similar to pipelineiterator in its unpacking mechanism but we have an extra required item which is the presence of is_last that is because everything is flattened by pipelinechunkiterator we need to keep track of how to regroup here in the original process boundaries so that process and postprocess see the same data this iterator accumulates items possibly while unbatching until it its a is_last and then just passes it on to the caller could be last batch so we can t unroll as many elements
import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput


class PipelineDataset(Dataset):
    def __init__(self, dataset, process, params):
        self.dataset = dataset
        self.process = process
        self.params = params

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, i):
        item = self.dataset[i]
        processed = self.process(item, **self.params)
        return processed


class PipelineIterator(IterableDataset):
    def __init__(self, loader, infer, params, loader_batch_size=None):
        
        self.loader = loader
        self.infer = infer
        self.params = params
        if loader_batch_size == 1:
            
            loader_batch_size = None
        self.loader_batch_size = loader_batch_size

        
        self._loader_batch_index = None
        self._loader_batch_data = None

    def __len__(self):
        return len(self.loader)

    def __iter__(self):
        self.iterator = iter(self.loader)
        return self

    def loader_batch_item(self):
        
        if isinstance(self._loader_batch_data, torch.Tensor):
            
            result = self._loader_batch_data[self._loader_batch_index]
        else:
            
            loader_batched = {}
            for k, element in self._loader_batch_data.items():
                if isinstance(element, ModelOutput):
                    
                    element = element.to_tuple()
                    if isinstance(element[0], torch.Tensor):
                        loader_batched[k] = tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
                    elif isinstance(element[0], np.ndarray):
                        loader_batched[k] = tuple(np.expand_dims(el[self._loader_batch_index], 0) for el in element)
                    continue
                if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(element, tuple):
                    
                    if isinstance(element[0], torch.Tensor):
                        loader_batched[k] = tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
                    elif isinstance(element[0], np.ndarray):
                        loader_batched[k] = tuple(np.expand_dims(el[self._loader_batch_index], 0) for el in element)
                    continue
                if element is None:
                    
                    loader_batched[k] = None
                elif isinstance(element[self._loader_batch_index], torch.Tensor):
                    
                    

                    loader_batched[k] = element[self._loader_batch_index].unsqueeze(0)
                elif isinstance(element[self._loader_batch_index], np.ndarray):
                    
                    
                    loader_batched[k] = np.expand_dims(element[self._loader_batch_index], 0)
                else:
                    
                    loader_batched[k] = element[self._loader_batch_index]
            
            
            result = self._loader_batch_data.__class__(loader_batched)
        self._loader_batch_index += 1
        return result

    def __next__(self):
        if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
            
            
            return self.loader_batch_item()

        
        item = next(self.iterator)
        processed = self.infer(item, **self.params)
        
        if self.loader_batch_size is not None:
            
            if isinstance(processed, torch.Tensor):
                first_tensor = processed
            else:
                key = list(processed.keys())[0]
                first_tensor = processed[key]
            if isinstance(first_tensor, list):
                observed_batch_size = len(first_tensor)
            else:
                observed_batch_size = first_tensor.shape[0]
            if 0 < observed_batch_size < self.loader_batch_size:
                
                
                self.loader_batch_size = observed_batch_size
            
            self._loader_batch_data = processed
            self._loader_batch_index = 0
            return self.loader_batch_item()
        else:
            
            return processed


class PipelineChunkIterator(PipelineIterator):
    def __init__(self, loader, infer, params, loader_batch_size=None):
        
        super().__init__(loader, infer, params)

    def __iter__(self):
        self.iterator = iter(self.loader)
        self.subiterator = None
        return self

    def __next__(self):
        if self.subiterator is None:
            "Subiterator None means we haven't started a `preprocess` iterator. so start it"
            self.subiterator = self.infer(next(self.iterator), **self.params)
        try:
            
            processed = next(self.subiterator)
        except StopIteration:
            
            
            
            
            
            
            self.subiterator = self.infer(next(self.iterator), **self.params)
            processed = next(self.subiterator)
        return processed


class PipelinePackIterator(PipelineIterator):
    

    def __iter__(self):
        self.iterator = iter(self.loader)
        return self

    def __next__(self):
        
        
        
        
        

        
        
        is_last = False
        accumulator = []
        if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
            while self._loader_batch_index < self.loader_batch_size:
                item = self.loader_batch_item()
                is_last = item.pop("is_last")
                accumulator.append(item)
                if is_last:
                    return accumulator

        while not is_last:
            processed = self.infer(next(self.iterator), **self.params)
            if self.loader_batch_size is not None:
                if isinstance(processed, torch.Tensor):
                    first_tensor = processed
                else:
                    key = list(processed.keys())[0]
                    first_tensor = processed[key]
                if isinstance(first_tensor, list):
                    observed_batch_size = len(first_tensor)
                else:
                    observed_batch_size = first_tensor.shape[0]
                if 0 < observed_batch_size < self.loader_batch_size:
                    
                    
                    self.loader_batch_size = observed_batch_size
                self._loader_batch_data = processed
                self._loader_batch_index = 0
                while self._loader_batch_index < self.loader_batch_size:
                    item = self.loader_batch_item()
                    is_last = item.pop("is_last")
                    accumulator.append(item)
                    if is_last:
                        return accumulator
            else:
                item = processed
                is_last = item.pop("is_last")
                accumulator.append(item)
        return accumulator


class KeyDataset(Dataset):
    def __init__(self, dataset: Dataset, key: str):
        self.dataset = dataset
        self.key = key

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, i):
        return self.dataset[i][self.key]


class KeyPairDataset(Dataset):
    def __init__(self, dataset: Dataset, key1: str, key2: str):
        self.dataset = dataset
        self.key1 = key1
        self.key2 = key2

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, i):
        return {"text": self.dataset[i][self.key1], "text_pair": self.dataset[i][self.key2]}