Merging Two datasets

datasets.py

import tensorflow as tf

def approach_1():
  def gen():
    l = ["foo", "bar", "baz"]
    
    l1 = [f"{x}_py" for x in l]
    l2 = [f"{x}_java" for x in l]
    
    d1 = tf.data.Dataset.from_tensor_slices(l1)
    d2 = tf.data.Dataset.from_tensor_slices(l2)

    for data in iter(tf.data.Dataset.zip((d1,d2))):
      yield data[0]
      yield data[1]

  signature = tf.TensorSpec(shape=(), dtype=tf.string)
  ds = tf.data.Dataset.from_generator(gen, output_signature=signature).batch(4)
  for e in iter(ds):
    print(e)


def approach_2():
  """WIP"""
  l = ["foo", "bar", "baz"]

  l1 = [f"{x}_py" for x in l]
  l2 = [f"{x}_java" for x in l]

  d1 = tf.data.Dataset.from_tensor_slices(l1)
  d2 = tf.data.Dataset.from_tensor_slices(l2)

  ds = tf.data.Dataset.zip((d1, d2))
  ds.flat_map(lambda x, y: tf.data.Dataset.from_tensor_slices((x, y)))
  print(next(iter(ds)))

if __name__ == "__main__":
  approach_1()
  # approach_2()