Repository: OpenNMT/OpenNMT-tf

Issue:
An issue with SequenceRecordInputter ?
I tried to create SequenceClassifier model which used SequenceRecordInputter as a part of ParallelInputter - it produced an error before ending first learning step. After isolating the problem, it seems that SequenceRecordInputter dataset generation is the source of it:

Reproducible code:
```python3
import numpy as np
from opennmt import encoders, inputters, models, Runner

vectors = []
for i in range(1000):
    vectors.append(np.random.rand(np.random.randint(1, 9), 16))
inputters.create_sequence_records(vectors, "train.records")
with open("train_labels.txt", "w") as f:
    f.write("\n".join(np.random.randint(0, 2, 1000).astype("str")))
with open("labels_vocab.txt", "w") as f:
    f.write("\n".join(["0", "1"]))
model = models.SequenceClassifier(
    inputters.SequenceRecordInputter(16),
    encoders.SelfAttentionEncoder(
        num_layers=2, num_units=16, num_heads=4, ffn_inner_dim=64
    ),
)
config = {
    "model_dir": ".",
    "data": {
        "target_vocabulary": "labels_vocab.txt",
        "train_features_file": "train.records",
        "train_labels_file": "train_labels.txt",
    },
    "params": {"optimizer": "Adam", "learning_rate": 0.001},
    "train": {"batch_size": 1, "max_step": 2},
}
runner = Runner(model, config, auto_config=False)
runner.train()
```

Error text
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [12], in <cell line: 1>()
----> 1 runner.train()

File ~/.local/lib/python3.9/site-packages/opennmt/runner.py:281, in Runner.train(self, num_devices, with_eval, checkpoint_path, hvd, return_summary, fallback_to_cpu, continue_from_checkpoint)
    278 else:
    279     trainer = training_util.Trainer(model, optimizer, checkpoint=checkpoint)
--> 281 summary = trainer(
    282     dataset_fn,
    283     max_step=train_config.get("max_step"),
    284     accum_steps=accum_steps,
    285     report_steps=train_config.get("save_summary_steps", 100),
    286     save_steps=train_config.get("save_checkpoints_steps", 5000),
    287     evaluator=evaluator,
    288     eval_steps=eval_config.get("steps", 5000),
    289     moving_average_decay=train_config.get("moving_average_decay"),
    290 )
    292 average_last_checkpoints = train_config.get("average_last_checkpoints", 0)
    293 if checkpoint is None:

File ~/.local/lib/python3.9/site-packages/opennmt/training.py:109, in Trainer.__call__(self, dataset, max_step, accum_steps, report_steps, save_steps, evaluator, eval_steps, moving_average_decay)
    107 step = None
    108 moving_average = None
--> 109 for i, loss in enumerate(
    110     self._steps(dataset, accum_steps=accum_steps, report_steps=report_steps)
    111 ):
    112     if i == 0:
    113         self._log_model_info()

File ~/.local/lib/python3.9/site-packages/opennmt/training.py:221, in Trainer._steps(self, dataset, accum_steps, report_steps)
    209 def _steps(self, dataset, accum_steps=1, report_steps=None):
    210     """Returns a generator over training steps (i.e. parameters update).
    211 
    212     Args:
   (...)
    219       A generator that yields a loss value to report for this step.
    220     """
--> 221     dataset = self._finalize_dataset(dataset)
    222     iterator = iter(dataset)
    224     # We define 2 separate functions to support gradient accumulation:
    225     #  * forward: compute and accumulate the gradients
    226     #  * step: apply the gradients
    227     # When gradient accumulation is disabled, the forward function also applies the gradients.

File ~/.local/lib/python3.9/site-packages/opennmt/training.py:206, in Trainer._finalize_dataset(self, dataset)
    196 """Returns the final dataset instance to be used for training.
    197 
    198 Args:
   (...)
    203   A ``tf.data.Dataset``.
    204 """
    205 if callable(dataset):
--> 206     dataset = dataset(tf.distribute.InputContext())
    207 return dataset

File ~/.local/lib/python3.9/site-packages/opennmt/runner.py:220, in Runner.train.<locals>.<lambda>(input_context)
    216 batch_type = train_config["batch_type"]
    217 batch_size_multiple = 8 if mixed_precision and batch_type == "tokens" else 1
    219 dataset_fn = (
--> 220     lambda input_context: model.examples_inputter.make_training_dataset(
    221         data_config["train_features_file"],
    222         data_config.get("train_labels_file"),
    223         train_config["batch_size"],
    224         batch_type=batch_type,
    225         batch_size_multiple=batch_size_multiple,
    226         shuffle_buffer_size=train_config["sample_buffer_size"],
    227         length_bucket_width=train_config["length_bucket_width"],
    228         maximum_features_length=train_config.get("maximum_features_length"),
    229         maximum_labels_length=train_config.get("maximum_labels_length"),
    230         single_pass=train_config.get("single_pass", False),
    231         num_shards=input_context.num_input_pipelines,
    232         shard_index=input_context.input_pipeline_id,
    233         prefetch_buffer_size=train_config.get("prefetch_buffer_size"),
    234         cardinality_multiple=input_context.num_replicas_in_sync,
    235         weights=data_config.get("train_files_weights"),
    236         batch_autotune_mode=train_config.get("batch_autotune_mode"),
    237     )
    238 )
    240 checkpoint = None
    241 evaluator = None

File ~/.local/lib/python3.9/site-packages/opennmt/inputters/inputter.py:834, in ExampleInputterAdapter.make_training_dataset(self, features_file, labels_file, batch_size, batch_type, batch_multiplier, batch_size_multiple, shuffle_buffer_size, length_bucket_width, maximum_features_length, maximum_labels_length, single_pass, num_shards, shard_index, num_threads, prefetch_buffer_size, cardinality_multiple, weights, batch_autotune_mode)
    832 if weights is not None:
    833     dataset = (dataset, weights)
--> 834 dataset = dataset_util.training_pipeline(
    835     batch_size,
    836     batch_type=batch_type,
    837     batch_multiplier=batch_multiplier,
    838     batch_size_multiple=batch_size_multiple,
    839     transform_fns=transform_fns,
    840     length_bucket_width=length_bucket_width,
    841     features_length_fn=features_length_fn,
    842     labels_length_fn=labels_length_fn,
    843     single_pass=single_pass,
    844     num_shards=num_shards,
    845     shard_index=shard_index,
    846     num_threads=num_threads,
    847     dataset_size=self.get_dataset_size(data_files),
    848     shuffle_buffer_size=shuffle_buffer_size,
    849     prefetch_buffer_size=prefetch_buffer_size,
    850     cardinality_multiple=cardinality_multiple,
    851 )(dataset)
    852 return dataset

File ~/.local/lib/python3.9/site-packages/opennmt/data/dataset.py:637, in training_pipeline.<locals>._pipeline(dataset)
    635 if labels_length_fn is not None:
    636     length_fn.append(labels_length_fn)
--> 637 dataset = dataset.apply(
    638     batch_sequence_dataset(
    639         batch_size,
    640         batch_type=batch_type,
    641         batch_multiplier=batch_multiplier,
    642         batch_size_multiple=batch_size_multiple,
    643         length_bucket_width=length_bucket_width,
    644         length_fn=length_fn,
    645     )
    646 )
    647 dataset = dataset.apply(filter_irregular_batches(batch_multiplier))
    648 if not single_pass:

File ~/.local/lib/python3.9/site-packages/tensorflow/python/data/ops/dataset_ops.py:2270, in DatasetV2.apply(self, transformation_func)
   2248 def apply(self, transformation_func):
   2249   """Applies a transformation function to this dataset.
   2250 
   2251   `apply` enables chaining of custom `Dataset` transformations, which are
   (...)
   2268         dataset.
   2269   """
-> 2270   dataset = transformation_func(self)
   2271   if not isinstance(dataset, DatasetV2):
   2272     raise TypeError(
   2273         f"`transformation_func` must return a `tf.data.Dataset` object. "
   2274         f"Got {type(dataset)}.")

File ~/.local/lib/python3.9/site-packages/opennmt/data/dataset.py:482, in batch_sequence_dataset.<locals>.<lambda>(dataset)
    475 else:
    476     raise ValueError(
    477         "Invalid batch type: '{}'; should be 'examples' or 'tokens'".format(
    478             batch_type
    479         )
    480     )
--> 482 return lambda dataset: dataset.group_by_window(_key_func, _reduce_func, **kwargs)

File ~/.local/lib/python3.9/site-packages/tensorflow/python/data/ops/dataset_ops.py:2823, in DatasetV2.group_by_window(self, key_func, reduce_func, window_size, window_size_func, name)
   2819   window_size_func = constant_window_func
   2821 assert window_size_func is not None
-> 2823 return _GroupByWindowDataset(
   2824     self, key_func, reduce_func, window_size_func, name=name)

File ~/.local/lib/python3.9/site-packages/tensorflow/python/data/ops/dataset_ops.py:5683, in _GroupByWindowDataset.__init__(self, input_dataset, key_func, reduce_func, window_size_func, name)
   5681 """See `group_by_window()` for details."""
   5682 self._input_dataset = input_dataset
-> 5683 self._make_key_func(key_func, input_dataset)
   5684 self._make_reduce_func(reduce_func, input_dataset)
   5685 self._make_window_size_func(window_size_func)

File ~/.local/lib/python3.9/site-packages/tensorflow/python/data/ops/dataset_ops.py:5721, in _GroupByWindowDataset._make_key_func(self, key_func, input_dataset)
   5718 def key_func_wrapper(*args):
   5719   return ops.convert_to_tensor(key_func(*args), dtype=dtypes.int64)
-> 5721 self._key_func = structured_function.StructuredFunctionWrapper(
   5722     key_func_wrapper, self._transformation_name(), dataset=input_dataset)
   5723 if not self._key_func.output_structure.is_compatible_with(
   5724     tensor_spec.TensorSpec([], dtypes.int64)):
   5725   raise ValueError(f"Invalid `key_func`. `key_func` must return a single "
   5726                    f"`tf.int64` scalar tensor but its return type is "
   5727                    f"{self._key_func.output_structure}.")

File ~/.local/lib/python3.9/site-packages/tensorflow/python/data/ops/structured_function.py:271, in StructuredFunctionWrapper.__init__(self, func, transformation_name, dataset, input_classes, input_shapes, input_types, input_structure, add_to_graph, use_legacy_function, defun_kwargs)
    264       warnings.warn(
    265           "Even though the `tf.config.experimental_run_functions_eagerly` "
    266           "option is set, this option does not apply to tf.data functions. "
    267           "To force eager execution of tf.data functions, please use "
    268           "`tf.data.experimental.enable_debug_mode()`.")
    269     fn_factory = trace_tf_function(defun_kwargs)
--> 271 self._function = fn_factory()
    272 # There is no graph to add in eager mode.
    273 add_to_graph &= not context.executing_eagerly()

File ~/.local/lib/python3.9/site-packages/tensorflow/python/eager/function.py:2567, in Function.get_concrete_function(self, *args, **kwargs)
   2558 def get_concrete_function(self, *args, **kwargs):
   2559   """Returns a `ConcreteFunction` specialized to inputs and execution context.
   2560 
   2561   Args:
   (...)
   2565        or `tf.Tensor` or `tf.TensorSpec`.
   2566   """
-> 2567   graph_function = self._get_concrete_function_garbage_collected(
   2568       *args, **kwargs)
   2569   graph_function._garbage_collector.release()  # pylint: disable=protected-access
   2570   return graph_function

File ~/.local/lib/python3.9/site-packages/tensorflow/python/eager/function.py:2533, in Function._get_concrete_function_garbage_collected(self, *args, **kwargs)
   2531   args, kwargs = None, None
   2532 with self._lock:
-> 2533   graph_function, _ = self._maybe_define_function(args, kwargs)
   2534   seen_names = set()
   2535   captured = object_identity.ObjectIdentitySet(
   2536       graph_function.graph.internal_captures)

File ~/.local/lib/python3.9/site-packages/tensorflow/python/eager/function.py:2711, in Function._maybe_define_function(self, args, kwargs)
   2708   cache_key = self._function_cache.generalize(cache_key)
   2709   (args, kwargs) = cache_key._placeholder_value()  # pylint: disable=protected-access
-> 2711 graph_function = self._create_graph_function(args, kwargs)
   2712 self._function_cache.add(cache_key, cache_key_deletion_observer,
   2713                          graph_function)
   2715 return graph_function, filtered_flat_args

File ~/.local/lib/python3.9/site-packages/tensorflow/python/eager/function.py:2627, in Function._create_graph_function(self, args, kwargs)
   2622 missing_arg_names = [
   2623     "%s_%d" % (arg, i) for i, arg in enumerate(missing_arg_names)
   2624 ]
   2625 arg_names = base_arg_names + missing_arg_names
   2626 graph_function = ConcreteFunction(
-> 2627     func_graph_module.func_graph_from_py_func(
   2628         self._name,
   2629         self._python_function,
   2630         args,
   2631         kwargs,
   2632         self.input_signature,
   2633         autograph=self._autograph,
   2634         autograph_options=self._autograph_options,
   2635         arg_names=arg_names,
   2636         capture_by_value=self._capture_by_value),
   2637     self._function_attributes,
   2638     spec=self.function_spec,
   2639     # Tell the ConcreteFunction to clean up its graph once it goes out of
   2640     # scope. This is not the default behavior since it gets used in some
   2641     # places (like Keras) where the FuncGraph lives longer than the
   2642     # ConcreteFunction.
   2643     shared_func_graph=False)
   2644 return graph_function

File ~/.local/lib/python3.9/site-packages/tensorflow/python/framework/func_graph.py:1141, in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, acd_record_initial_resource_uses)
   1138 else:
   1139   _, original_func = tf_decorator.unwrap(python_func)
-> 1141 func_outputs = python_func(*func_args, **func_kwargs)
   1143 # invariant: `func_outputs` contains only Tensors, CompositeTensors,
   1144 # TensorArrays and `None`s.
   1145 func_outputs = nest.map_structure(
   1146     convert, func_outputs, expand_composites=True)

File ~/.local/lib/python3.9/site-packages/tensorflow/python/data/ops/structured_function.py:248, in StructuredFunctionWrapper.__init__.<locals>.trace_tf_function.<locals>.wrapped_fn(*args)
    242 @eager_function.defun_with_attributes(
    243     input_signature=structure.get_flat_tensor_specs(
    244         self._input_structure),
    245     autograph=False,
    246     attributes=defun_kwargs)
    247 def wrapped_fn(*args):  # pylint: disable=missing-docstring
--> 248   ret = wrapper_helper(*args)
    249   ret = structure.to_tensor_list(self._output_structure, ret)
    250   return [ops.convert_to_tensor(t) for t in ret]

File ~/.local/lib/python3.9/site-packages/tensorflow/python/data/ops/structured_function.py:177, in StructuredFunctionWrapper.__init__.<locals>.wrapper_helper(*args)
    175 if not _should_unpack(nested_args):
    176   nested_args = (nested_args,)
--> 177 ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
    178 if _should_pack(ret):
    179   ret = tuple(ret)

File ~/.local/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py:689, in convert.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
    687 try:
    688   with conversion_ctx:
--> 689     return converted_call(f, args, kwargs, options=options)
    690 except Exception as e:  # pylint:disable=broad-except
    691   if hasattr(e, 'ag_error_metadata'):

File ~/.local/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py:377, in converted_call(f, args, kwargs, caller_fn_scope, options)
    374   return _call_unconverted(f, args, kwargs, options)
    376 if not options.user_requested and conversion.is_allowlisted(f):
--> 377   return _call_unconverted(f, args, kwargs, options)
    379 # internal_convert_user_code is for example turned off when issuing a dynamic
    380 # call conversion from generated code while in nonrecursive mode. In that
    381 # case we evidently don't want to recurse, but we still have to convert
    382 # things like builtins.
    383 if not options.internal_convert_user_code:

File ~/.local/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py:458, in _call_unconverted(f, args, kwargs, options, update_cache)
    455   return f.__self__.call(args, kwargs)
    457 if kwargs is not None:
--> 458   return f(*args, **kwargs)
    459 return f(*args)

File ~/.local/lib/python3.9/site-packages/tensorflow/python/data/ops/dataset_ops.py:5719, in _GroupByWindowDataset._make_key_func.<locals>.key_func_wrapper(*args)
   5718 def key_func_wrapper(*args):
-> 5719   return ops.convert_to_tensor(key_func(*args), dtype=dtypes.int64)

File ~/.local/lib/python3.9/site-packages/opennmt/data/dataset.py:442, in batch_sequence_dataset.<locals>._key_func(*args)
    437     raise ValueError(
    438         "%d length functions were passed but this dataset contains "
    439         "%d parallel elements" % (len(length_fns), len(args))
    440     )
    441 # Take the highest bucket id.
--> 442 bucket_id = tf.reduce_max(
    443     [
    444         _get_bucket_id(features, length_fn)
    445         for features, length_fn in zip(args, length_fns)
    446     ]
    447 )
    448 return tf.cast(bucket_id, tf.int64)

File ~/.local/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File ~/.local/lib/python3.9/site-packages/tensorflow/python/ops/array_ops.py:1506, in _autopacking_helper(list_or_tuple, dtype, name)
   1504 if isinstance(elem, core.Tensor):
   1505   if dtype is not None and elem.dtype.base_dtype != dtype:
-> 1506     raise TypeError(f"Cannot convert a list containing a tensor of dtype "
   1507                     f"{elem.dtype} to {dtype} (Tensor is: {elem!r})")
   1508   converted_elems.append(elem)
   1509   must_pack = True

TypeError: Cannot convert a list containing a tensor of dtype <dtype: 'int32'> to <dtype: 'int64'> (Tensor is: <tf.Tensor 'Const_1:0' shape=() dtype=int32>)
```


Discussion:

Resolving patch:
diff --git a/opennmt/inputters/record_inputter.py b/opennmt/inputters/record_inputter.py
--- a/opennmt/inputters/record_inputter.py
+++ b/opennmt/inputters/record_inputter.py
@@ -48,7 +48,7 @@ def make_features(self, element=None, features=None, training=None):
             },
         )
         tensor = feature_lists["values"]
-        features["length"] = lengths["values"]
+        features["length"] = tf.cast(lengths["values"], tf.int32)
         features["tensor"] = tf.cast(tensor, self.dtype)
         return features
 


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -770,6 +770,7 @@ def testSequenceRecordBatch(self):
         features = next(iter(dataset))
         lengths = features["length"]
         tensors = features["tensor"]
+        self.assertEqual(lengths.dtype, tf.int32)
         self.assertAllEqual(lengths, [3, 6, 1])
         for length, tensor, expected_vector in zip(lengths, tensors, vectors):
             self.assertAllClose(tensor[:length], expected_vector)