Fail to run trainer.train() with huggingface transformer - tensorflow

Question

I am trying to set up a TensorFlow fine-tune framework for a question-answering project. Using hugging-face/transformer as the prototype, but cannot run through the trainer.
The experiment is conducted at Databricks, the pre-trained model loaded is base-bert, train and dev sets are downloaded from hugging-face examples SQUAD 2.0 https://github.com/huggingface/transformers/tree/master/examples/question-answering
The error log complains about the unexpected keyword argument 'is_impossible', which is a SQUAD 2 data format feature.
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
model = TFAutoModelForQuestionAnswering.from_pretrained('bert-base-uncased')
processor = SquadV2Processor()
train_examples = processor.get_train_examples(data_dir, train_file_name)
eval_examples = processor.get_dev_examples(data_dir, dev_file_name)
train_dataset = (squad_convert_examples_to_features(
examples=train_examples,
tokenizer=tokenizer,
max_seq_length=max_seq_length,
doc_stride=doc_stride,
max_query_length=max_query_length,
is_training=True,
return_dataset="tf"
))
eval_dataset = (squad_convert_examples_to_features(
examples=eval_examples,
tokenizer=tokenizer,
max_seq_length=max_seq_length,
doc_stride=doc_stride,
max_query_length=max_query_length,
is_training=False,
return_dataset="tf"
))
training_args = (TFTrainingArguments(
output_dir=output_dir,
num_train_epochs=2,
do_train=True,
per_device_train_batch_size = 8,
per_device_eval_batch_size = 16,
logging_steps=10,
learning_rate=3e-5))
trainer = TFTrainer(model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset)
trainer.train()
Error Log as below:
TypeError Traceback (most recent call
last) in
----> 1 trainer.train()
/databricks/python/lib/python3.7/site-packages/transformers/trainer_tf.py
in train(self)
410 if self.args.past_index >= 0:
411 self._past = None
--> 412 for step, training_loss in enumerate(self._training_steps(train_ds, optimizer)):
413 self.global_step = iterations.numpy()
414 self.epoch_logging = epoch_iter - 1 + (step + 1) / steps_per_epoch
/databricks/python/lib/python3.7/site-packages/transformers/trainer_tf.py
in _training_steps(self, ds, optimizer)
457 Returns a generator over training steps (i.e. parameters update).
458 """
--> 459 for i, loss in enumerate(self._accumulate_next_gradients(ds)):
460 if i % self.args.gradient_accumulation_steps == 0:
461 self._apply_gradients(optimizer)
/databricks/python/lib/python3.7/site-packages/transformers/trainer_tf.py
in _accumulate_next_gradients(self, ds)
490 while True:
491 try:
--> 492 yield _accumulate_next()
493 except tf.errors.OutOfRangeError:
494 break
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py
in call(self, *args, **kwds)
566 xla_context.Exit()
567 else:
--> 568 result = self._call(*args, **kwds)
569
570 if tracing_count == self._get_tracing_count():
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py
in _call(self, *args, **kwds)
613 # This is the first call of call, so we have to initialize.
614 initializers = []
--> 615 self._initialize(args, kwds, add_initializers_to=initializers)
616 finally:
617 # At this point we know that the initialization is complete (or less
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py
in _initialize(self, args, kwds, add_initializers_to)
495 self._concrete_stateful_fn = (
496 self._stateful_fn._get_concrete_function_internal_garbage_collected(
pylint: disable=protected-access
--> 497 *args, **kwds))
498
499 def invalid_creator_scope(*unused_args, **unused_kwds):
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py
in _get_concrete_function_internal_garbage_collected(self, *args,
**kwargs)
2387 args, kwargs = None, None
2388 with self._lock:
-> 2389 graph_function, _, _ = self._maybe_define_function(args, kwargs)
2390 return graph_function
2391
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py
in _maybe_define_function(self, args, kwargs)
2701
2702 self._function_cache.missed.add(call_context_key)
-> 2703 graph_function = self._create_graph_function(args, kwargs)
2704 self._function_cache.primary[cache_key] = graph_function
2705 return graph_function, args, kwargs
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py
in _create_graph_function(self, args, kwargs,
override_flat_arg_shapes)
2591 arg_names=arg_names,
2592 override_flat_arg_shapes=override_flat_arg_shapes,
-> 2593 capture_by_value=self._capture_by_value),
2594 self._function_attributes,
2595 # Tell the ConcreteFunction to clean up its graph once it goes out of
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py
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,
override_flat_arg_shapes)
976 converted_func)
977
--> 978 func_outputs = python_func(*func_args, **func_kwargs)
979
980 # invariant: func_outputs contains only Tensors, CompositeTensors,
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py
in wrapped_fn(*args, **kwds)
437 # wrapped allows AutoGraph to swap in a converted function. We give
438 # the function a weak reference to itself to avoid a reference cycle.
--> 439 return weak_wrapped_fn().wrapped(*args, **kwds)
440 weak_wrapped_fn = weakref.ref(wrapped_fn)
441
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py
in wrapper(*args, **kwargs)
966 except Exception as e: # pylint:disable=broad-except
967 if hasattr(e, "ag_error_metadata"):
--> 968 raise e.ag_error_metadata.to_exception(e)
969 else:
970 raise
TypeError: in converted code:
/databricks/python/lib/python3.7/site-packages/transformers/trainer_tf.py:488
_accumulate_next *
return self._accumulate_gradients(per_replica_features, per_replica_labels)
/databricks/python/lib/python3.7/site-packages/transformers/trainer_tf.py:498
_accumulate_gradients *
per_replica_loss = self.args.strategy.experimental_run_v2(
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/distribute/one_device_strategy.py:180 experimental_run_v2
return super(OneDeviceStrategy, self).experimental_run_v2(fn, args, kwargs)
/databricks/python/lib/python3.7/site-packages/transformers/trainer_tf.py:511
_forward *
per_example_loss, _ = self._run_model(features, labels, True)
/databricks/python/lib/python3.7/site-packages/transformers/trainer_tf.py:532
_run_model *
outputs = self.model(features, training=training, **labels)[:2]
/local_disk0/pythonVirtualEnvDirs/virtualEnv-f56565e5-e45b-447a-b7df-50daf9109495/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py:778
call
outputs = call_fn(cast_inputs, *args, **kwargs)
TypeError: tf__call() got an unexpected keyword argument 'is_impossible'