cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Remove batch dimension from the output of tflite_with_tokenizer in text classifier.

Browse Source 
PiperOrigin-RevId: 581292824
This commit is contained in:
MediaPipe Team 2023-11-10 10:00:56 -08:00 committed by Copybara-Service
parent d772bf8134
commit 64b21d758e
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mediapipe/model_maker/python/text/text_classifier/model_with_tokenizer.py
View File

@ -32,4 +32,4 @@ class ModelWithTokenizer(tf.keras.Model):
x = self._tokenizer.process_fn(input_tensor) x = self._tokenizer.process_fn(input_tensor)
x = {k: tf.expand_dims(v, axis=0) for k, v in x.items()} x = {k: tf.expand_dims(v, axis=0) for k, v in x.items()}
x = self._model(x) x = self._model(x)
return x return x[0] # TODO: Add back the batch dimension

2
mediapipe/model_maker/python/text/text_classifier/model_with_tokenizer_test.py
View File

@ -97,7 +97,7 @@ class BertTokenizerTest(tf.test.TestCase):
self._tokenizer, self._model self._tokenizer, self._model
) )
output = model(tf.constant(["Example input".encode("utf-8")])) output = model(tf.constant(["Example input".encode("utf-8")]))
self.assertAllEqual(output.shape, (1, 2)) self.assertAllEqual(output.shape, (2,))
self.assertEqual(tf.reduce_sum(output), 1) self.assertEqual(tf.reduce_sum(output), 1)