Handle ColumnSchema target in serialization of SequenceTransform

[oliverholworthy]

Fixes #1212

Goals ⚽

Enable serialization of SequenceTransform layers. with schema target.

Implementation Details 🚧

Handle ColumnSchema target in serialization of SequenceTransform

Testing Details 🔍

[github-actions]

Documentation preview

https://nvidia-merlin.github.io/models/review/pr-1214

[oliverholworthy]

While this change will enable saving a model with SequenceTransform layers with a target type of ColumnSchema.

The example code linked in the related bug issue, shows that the SequenceMaskLast layer is only used as part of the pre of the .evaluate call. This layer then fails to serialize when calling .save() on the model, because keras finds a ColumnSchema object which it doesn't know what to do with.

We should check if it's intended that the call to evaluate has a side-effect of adding this SequenceMaskLast layer to the model. Since if the SequenceTransform layers are only used in the pre, and don't become part of the model, in that usecase, they may not strictly need to be saved/serialized to json correctly.

[rnyak]

@gabrielspmoreira @sararb fyi.

[rnyak]

@gabrielspmoreira and @oliverholworthy I tested this fix and now we can serve topk_model on Triton, however served topK_model returns only top 100 scores which does not make sense without item ids. we might need to fix the pertinent code so that it returns both the scores and corresponding item_ids. I assume for that we need exported schema to show to outputs (scores and ids) ?