Creating SequenceExample(s) using TensorFlow Transform

Question

With TensorFlow Transform, we can pre-process data using Apache Beam. One of the requirements when setting up such a pipeline is to define a DatasetMetadata object, which contains the schema that has the information needed to parse the data from its on-disk or in-memory format, into tensors.

In the official documentation, we are given an example of the form:

raw_data_metadata = dataset_metadata.DatasetMetadata(
dataset_schema.from_feature_spec({
    's': tf.FixedLenFeature([], tf.string),
    'y': tf.FixedLenFeature([], tf.float32),
    'x': tf.FixedLenFeature([], tf.float32),
}))

This is all fine if your raw data is a dictionary of the form:

{
    's': 'example string',
    'y': 32.0,
    'x': 35.0
}

However, I am somewhat lost when it comes to defining a schema for a SequenceExample. More specifically, consider that my data has the following format:

{
    # context features
    'length': 5,
    # sequence features
    'tokens': [
        {
            'raw': 'The',
            'ner-tag': 'O'
        },
        {
            'raw': 'European',
            'ner-tag': 'B-org'
        },
        {
            'raw': 'Union',
            'ner-tag': 'I-org'
        },
        {
            'raw': 'is',
            'ner-tag': 'O'
        },
        {
            'raw': 'nice',
            'ner-tag': 'O'
        }
        ...
    ]
}

Above I have a sentence with 2 sequences:

ner-tag sequence which is going to be used as a label for the model
raw sequence which is going to be used as a feature for the model

How can I create a TFT data schema for such examples?

The documentation is a bit absent for this one. Any help much appreciated!

Andrei Balici Andrei Balici · Accepted Answer · 2020-01-12T13:06:31

Well, after more research, the answer is that you can't.

TensorFlow Transform doesn't yet support SequenceExample(s). Check this.

It appears that the only way to do this currently is to have the Beam Pipeline create the SequenceExamples, serialize them and write them to TFRecords.

Given the above sentence object structure, you'll need to create a Beam DoFn first that converts each sentence to a serialized SequenceExample:

class ConvertJSONSentenceToSerializedSequenceExample(beam.DoFn):

    def make_example(self, sentence):
        # the context features
        sentence_level_details = tf.train.Features(feature={
            'length': tf.train.Feature(int64_list=tf.train.Int64List(value=[sentence['length']]))
        })

        # create sequence data
        word_features = []
        ner_tags_features = []
        for token in sentence['tokens']:
            # create each of the features, then add them to the corresponding feature list
            word_feature = tf.train.Feature(bytes_list=tf.train.BytesList(value=[token['raw'].encode('utf-8')]))
            word_features.append(word_feature)

            ner_tag_feature = tf.train.Feature(int64_list=tf.train.Int64List(value=[token['']]))
            ner_tags_features.append(ner_tag_feature)

        words = tf.train.FeatureList(feature=word_features)
        ner_tags = tf.train.FeatureList(feature=ner_tags_features)

        sentence_sequences = tf.train.FeatureLists(feature_list={
            'words': words,
            'ner-tags': ner_tags
        })

        ex = tf.train.SequenceExample(
            context = sentence_level_details,
            feature_lists = sentence_sequences
        )

        return ex

    def process(self, sentence, **kwargs):
        try:
            ex = self.make_example(sentence)
            yield ex.SerializeToString()
        except Exception as e:
            logging.warning("JSON sentence could not be converted into SequenceExample: " + str(e))
            return None

Once this is done, you can then use the beam.io.tfrecordio module these serialized SequenceExample(s) into TFRecord(s) files:

with beam.Pipeline(RUNNER, options=opts) as p:
(p
...
| 'Convert sentences to serialized TensorFlow SequenceExamples' >> beam.ParDo(ConvertJSONSentenceToSerializedSequenceExample())
| 'Write to TFRecord files' >> tfrecordio.WriteToTFRecord(
     os.path.join(OUTPUT_DIR, 'train'),
     file_name_suffix='.gz'
     # default coder is the BytesCoder, which will work since we have serialized the training data
)

Creating SequenceExample(s) using TensorFlow Transform

1 Answers