search

kyzhouhzau / BERT-NER

Use Google's BERT for named entity recognition （CoNLL-2003 as the dataset）.
MIT License
1.23k stars 335 forks source link

TPUEstimatorSpec.predictions must be dict of Tensors #19

Closed rikhuijzer closed 2 years ago

rikhuijzer commented 5 years ago

When running predict on Google Colab (to use TPU) the code crashes with the following error:

TPUEstimatorSpec.predictions must be dict of Tensors.

To solve it one can place the following code in _createmodel

predict = tf.argmax(probabilities, axis=-1)
predict_dict = {'predictions': predict}  # this way it is not shot down by check in TPUEstimatorSpec
return loss, per_example_loss, logits, predict_dict

This of course also means changing the interpretation of the result

result = estimator.predict(input_fn=predict_input_fn)
result = list(result)
result = [pred['predictions'] for pred in result]

Currently I'm unable to to pull request since that would mean looking into whether it really is a solution. Just posting it here for anyone who has the same problem.

amankhandelia commented 5 years ago

Thanks for the info

jloutz commented 5 years ago

awesome fix thanx - helped me trying to run biobert on google tpu