data-generator_2.py

def get_all(*args, chunksize):
  pds = []
  args = args[0]
  for arg in args:
    pds.append(pd.read_csv(f'{BASE_PATH}/{arg}', chunksize=chunksize))
  return pds
  
def merge_all(*args):
  merged = None
  args = args[0]
  for arg in args:
    if merged is None:
      tmp = arg
    else:
      merged = pd.merge(tmp, arg, how='inner', on=['protein', 'index'])
  return merged

def data_generator():
  total_row = NUMBER_OF_ROWS
  files = [
    'a.csv',
    'b.csv',
    'c.csv',
    'd.csv',
  ]

  pds = get_all(files, chunksize=batch_size)
  cnt = 0

  while True:
    data_frames = []
    for reader in pds:
      data_frames.append(reader.get_chunk())

    cnt += batch_size
    
    merged = merge_all(data_frames)
    x = merged.iloc[:, 1:].to_numpy()
    y = merged.iloc[:, 0].to_numpy()

    if cnt >= total_row:
      pds = get_all(files, chunksize=batch_size)
      cnt = 0

    yield x, y