code: def sentence_dataset_dict(
      filename,
      vocabs,
      attributes_input,
      attributes_target,
      batch_size,
      bucket_size,
      repeat=None,
      prefetch_size=tf.data.experimental.AUTOTUNE,
  ):
      """Combines sentences into a dataset of padded batches.
  
      Args:
        filename: file name of a corpus.
        vocabs: dictionary of dictionaries to map from strings to ids.
        attributes_input: attributes for the input.
        attributes_target: target attributes empty targets is not included.
        batch_size: the size of a batch.
        bucket_size: the size of a bucket.
        repeat: number of times the dataset is repeated.
        prefetch_size: prefetch size of the data.
  
      Returns:
        Returns dataset as dictionary containing the data as key value pairs.
      """
      data_keys = ["inputs"]
      if attributes_target:
          data_keys.append("targets")
  
      def generator():
          """Generator to create the data."""
          input_generator = sentences_from_conll_data(
              filename,
              vocabs,
              attributes_input,
              max_sentence_length=bucket_size,
          )
  
          if attributes_target:
              target_generator = sentences_from_conll_data(
                  filename,
                  vocabs,
                  attributes_target,
                  max_sentence_length=bucket_size,
              )
  
          for inputs in input_generator:
              data = {"inputs": inputs}
              if attributes_target:
                  data["targets"] = next(target_generator)
              yield data
  
      output_types = {k: tf.float32 for k in data_keys}
      output_shapes = {k: (None,) for k in data_keys}
      dataset = tf.data.Dataset.from_generator(
          generator,
          output_types=output_types,
          output_shapes=output_shapes,
      )
  
      # cache the dataset in memory and repeat.
      dataset = dataset.cache()
      dataset = dataset.repeat(repeat)
  
      # static padding up to bucket size.
      padded_shapes = {k: [bucket_size] for k in data_keys}
      dataset = dataset.padded_batch(
          batch_size=batch_size, padded_shapes=(padded_shapes)
      )
  
      dataset = dataset.prefetch(prefetch_size)
      return dataset
  
file_path: flax/examples/nlp_seq/input_pipeline.py
imports: import codecs
  import collections
  import enum
  import tensorflow
label_desc: Use image node_id 5 for calling a global function with example usage: sentence_dataset_dict(filename, vocabs, attributes_input, attributes_target, batch_size, bucket_size, repeat, prefetch_size) and returns: dataset