def test_TFRobertaForTokenClassification(self):
from transformers import RobertaConfig, TFRobertaForTokenClassification
keras.backend.clear_session()
# pretrained_weights = 'roberta-base'
tokenizer_file = 'roberta_roberta-base.pickle'
tokenizer = self._get_tokenzier(tokenizer_file)
text, inputs, inputs_onnx = self._prepare_inputs(tokenizer)
config = RobertaConfig()
model = TFRobertaForTokenClassification(config)
predictions = model.predict(inputs)
onnx_model = keras2onnx.convert_keras(model, model.name)
self.assertTrue(run_onnx_runtime(onnx_model.graph.name, onnx_model, inputs_onnx, predictions, self.model_files))
def getroBERTaModel():
config.num_labels = 2
roBERTModel = TFRobertaForTokenClassification.from_pretrained(
PRETRAINED_MODEL, config=config)
optimizer = tf.keras.optimizers.Adam(learning_rate=2e-5,
epsilon=1e-08,
clipnorm=1.0)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
roBERTModel.compile(optimizer=optimizer, loss=loss)
return roBERTModel
def test_TFRobertaForTokenClassification(self):
from transformers import RobertaTokenizer, TFRobertaForTokenClassification
pretrained_weights = 'roberta-base'
tokenizer = RobertaTokenizer.from_pretrained(pretrained_weights)
text, inputs, inputs_onnx = self._prepare_inputs(tokenizer)
model = TFRobertaForTokenClassification.from_pretrained(
pretrained_weights)
predictions = model.predict(inputs)
onnx_model = keras2onnx.convert_keras(model, model.name)
self.assertTrue(
run_onnx_runtime(onnx_model.graph.name, onnx_model, inputs_onnx,
predictions, self.model_files))