b156eba3ba09229a10692cffdc26bb8604b9f651,api-examples/transformer_utils.py,MultiFileLoader,__iter__,#MultiFileLoader#,410

Before Change

        files_per_worker = len(files) // all_workers
        offset = self.rank * num_workers_per_node + node_worker_id
        start_idx = offset * files_per_worker
        end_idx = start_idx + files_per_worker if offset < all_workers - 1 else len(files)
        print(f"worker {node_worker_id} [{start_idx}:{end_idx}]")

        self.vectorizer.mxlen = self.nctx

After Change

        all_workers = (self.world_size * num_workers_per_node)
        offset = self.rank * num_workers_per_node + node_worker_id
        self.vectorizer.mxlen = self.nctx
        read_file_order = list(range(offset, len(files), all_workers))
        // If we have multiple files per worker, possibly shuffle the file read order
        if self.shuffle:
            read_file_order = np.random.permutation(read_file_order)

        for file_idx in read_file_order:
            file = files[file_idx]
            with open(file) as rf:
                lines = rf.readlines()
                if self.shuffle:
                    random.shuffle(lines)
                for l in lines:
                    response = self.process_line(l)
                    yield response

Instances