Function produce_data

join/Deepjoin/all-mpnet-base-v2/train_script.py:168–220

(args, queue, filepaths, dataset_indices)

def produce_data(args, queue, filepaths, dataset_indices):
    global_batch_size = args.batch_size*args.nprocs    #Global batch size
    size_per_dataset = int(global_batch_size / args.datasets_per_batch)    #How many datasets per batch
    num_same_dataset = int(size_per_dataset / args.batch_size)
    print("producer", "global_batch_size", global_batch_size)
    print("producer", "size_per_dataset", size_per_dataset)
    print("producer", "num_same_dataset", num_same_dataset)
    
    datasets = []
    for filepath in filepaths:
        if "reddit_" in filepath:       #Special dataset class for Reddit files
            data_obj = RedditDataset(filepath)
        else:
            data_obj = Dataset(filepath)
        datasets.append(iter(data_obj)) 
    
    # Store if dataset is in a 2 col or 3 col format
    num_cols = {idx: len(next(dataset)) for idx, dataset in enumerate(datasets)}

    while True:
        texts_in_batch = set()
        batch_format = None     #2 vs 3 col format for this batch
        
        #Add data from several sub datasets
        for _ in range(args.datasets_per_batch):
            valid_dataset = False   #Check that datasets have the same 2/3 col format
            while not valid_dataset:
                data_idx = random.choice(dataset_indices)
                if batch_format is None:
                    batch_format = num_cols[data_idx]
                    valid_dataset = True
                else:   #Check that this dataset has the same format
                    valid_dataset = (batch_format == num_cols[data_idx])
            
            #Get data from this dataset
            dataset = datasets[data_idx]
            for _ in range(num_same_dataset):
                for _ in range(args.nprocs):
                    batch_device = []   #A batch for one device
                    while len(batch_device) < args.batch_size:
                        sample = next(dataset)
                        in_batch = False
                        for text in sample:
                            if text in texts_in_batch:
                                in_batch = True
                                break
                        
                        if not in_batch:
                            for text in sample:
                                texts_in_batch.add(text)
                            batch_device.append(sample)

                    queue.put(batch_device)
                      

class RedditDataset: