add tflite support

guesslang/guess.py

@@ -164,6 +164,14 @@ def _is_reliable(probabilities: List[float]) -> bool:
         predicted_language_probability = max(probabilities)
         return predicted_language_probability > threshold
 
+    def export(self,
+               ckpt_path: str,
+               tflite: bool = False) -> None:
+        with TemporaryDirectory() as model_logs_dir:
+            estimator = model.build(model_logs_dir,
+                                    list(self._extension_map))
+            model.save(estimator, self._saved_model_dir, ckpt_path, tflite)
+
 
 class GuesslangError(Exception):
     """Guesslang exception class"""

guesslang/model.py

@@ -1,16 +1,19 @@
 """Machine learning model"""
 
 from copy import deepcopy
+import os
+import functools
 import logging
 from operator import itemgetter
 from pathlib import Path
 import shutil
 from tempfile import TemporaryDirectory
-from typing import List, Tuple, Dict, Any, Callable
+from typing import List, Tuple, Dict, Any, Callable, Optional
 
 import tensorflow as tf
 from tensorflow.estimator import ModeKeys, Estimator
 from tensorflow.python.training.tracking.tracking import AutoTrackable
+import tensorflow_text as text
 
 
 LOGGER = logging.getLogger(__name__)
@@ -98,15 +101,31 @@ def train(estimator: Estimator, data_root_dir: str, max_steps: int) -> Any:
     return training_metrics
 
 
-def save(estimator: Estimator, saved_model_dir: str) -> None:
+def save(estimator: Estimator,
+         saved_model_dir: str,
+         ckpt_path: Optional[str] = None,
+         is_tflite: bool = False) -> None:
     """Save a Tensorflow estimator"""
     with TemporaryDirectory() as temporary_model_base_dir:
         export_dir = estimator.export_saved_model(
-            temporary_model_base_dir, _serving_input_receiver_fn
+            temporary_model_base_dir,
+            functools.partial(_serving_input_receiver_fn,
+                              is_tflite=is_tflite),
+            checkpoint_path=ckpt_path
         )
-
         Path(saved_model_dir).mkdir(exist_ok=True)
         export_path = Path(export_dir.decode()).absolute()
+        if is_tflite:
+            converter = tf.lite.TFLiteConverter.from_saved_model(
+                export_dir, signature_keys=['predict'])
+            converter.optimizations = [tf.lite.Optimize.DEFAULT]
+            converter.inference_type = tf.float32
+            ops_sets = [tf.lite.OpsSet.TFLITE_BUILTINS,
+                        tf.lite.OpsSet.SELECT_TF_OPS]
+            converter.target_spec.supported_ops = ops_sets
+            model = converter.convert()
+            tflite_path = os.path.join(saved_model_dir, 'guesslang.tflite')
+            tf.io.write_file(tflite_path, model)
         for path in export_path.glob('*'):
             shutil.move(str(path), saved_model_dir)
 
@@ -173,17 +192,24 @@ def input_function() -> tf.data.Dataset:
         if mode == ModeKeys.TRAIN:
             dataset = dataset.shuffle(Training.SHUFFLE_BUFFER).repeat()
 
-        return dataset.map(_preprocess).batch(HyperParameter.BATCH_SIZE)
+        return dataset.batch(HyperParameter.BATCH_SIZE).map(_preprocess)
 
     return input_function
 
 
-def _serving_input_receiver_fn() -> tf.estimator.export.ServingInputReceiver:
+def _serving_input_receiver_fn(
+        is_tflite: bool = False) -> tf.estimator.export.ServingInputReceiver:
     """Function to serve model for predictions."""
-
-    content = tf.compat.v1.placeholder(tf.string, [None])
-    receiver_tensors = {'content': content}
-    features = {'content': tf.map_fn(_preprocess_text, content)}
+    if is_tflite:
+        content = tf.compat.v1.placeholder(tf.string,
+                                           [HyperParameter.NB_TOKENS+1])
+        length = tf.compat.v1.placeholder(tf.int32, [])
+        receiver_tensors = {'content': content, 'length': length}
+        features = {'content': _preprocess_text_tflite(content, length)}
+    else:
+        content = tf.compat.v1.placeholder(tf.string, [None])
+        receiver_tensors = {'content': content}
+        features = {'content': _preprocess_text(content)}
 
     return tf.estimator.export.ServingInputReceiver(
         receiver_tensors=receiver_tensors,
@@ -209,9 +235,17 @@ def _preprocess(
 
 def _preprocess_text(data: tf.Tensor) -> tf.Tensor:
     """Feature engineering"""
-    padding = tf.constant(['']*HyperParameter.NB_TOKENS)
     data = tf.strings.bytes_split(data)
-    data = tf.strings.ngrams(data, HyperParameter.N_GRAM)
-    data = tf.concat((data, padding), axis=0)
-    data = data[:HyperParameter.NB_TOKENS]
-    return data
+    data = text.ngrams(
+        data, HyperParameter.N_GRAM, reduction_type=text.Reduction.STRING_JOIN)
+    return data.to_tensor(shape=(data.shape[0], HyperParameter.NB_TOKENS))
+
+
+def _preprocess_text_tflite(data: tf.Tensor, length: tf.Tensor) -> tf.Tensor:
+    processed_data, unprocessed_data = tf.split(
+        data, [length, HyperParameter.NB_TOKENS-length+1], num=2, axis=0)
+    processed_data = text.ngrams(
+        processed_data, HyperParameter.N_GRAM,
+        reduction_type=text.Reduction.STRING_JOIN)
+    return tf.expand_dims(
+        tf.concat([processed_data, unprocessed_data], axis=0), axis=0)

requirements.txt

@@ -1 +1,3 @@
-tensorflow==2.5.0
+tensorflow==2.6.0
+tensorflow-text==2.6.0
+keras==2.6.0