Support returning indices in text_dataset_from_directory

For data logging I need the indices ids of a row/ batch. Current implementation is limited in this regard:

tf.keras.utils.text_dataset_from_directory(
    directory
)
>> {"attention_mask" , "input_ids"}

For specific tasks the row id is needed:

tf.keras.utils.text_dataset_from_directory(
    directory,
    with_indices=True
)
>>  {"attention_mask" , "input_ids", "indices_ids"}

@haifeng-jin, @Haaris-Rahman, @mattdangerw, @guberti, @edumucelli

[chenmoneygithub]

yea I did not realize the map_fn has to be stateless. Another way is:

idx_ds = tf.data.experimental.Counter().take(ds.cardinality())
new_ds = tf.data.Dataset.zip((ds, idx_ds))

The new_ds object will have the index .

[franz101]

counter = -1
def fn(*data):
  global counter
  counter += 1
  return data, counter
a_iter = iter(train_ds.unbatch().map(fn))
a = next(a_iter)

trying this out in google colab, actually does not increment the counter