v 0.2.0 support tflite and serving

examples/convert_to_tflite.py

@@ -6,10 +6,12 @@
 from dialognlu import TransformerNLU
 from dialognlu.utils.tf_utils import convert_to_tflite_model
 
-model_path = "../saved_models/joint_distilbert_model"
+# model_path = "../saved_models/joint_distilbert_model"
+model_path = "../saved_models/joint_trans_bert_model"
 
 print("Loading model ...")
 nlu = TransformerNLU.load(model_path)
 
-save_file_path = "../saved_models/joint_distilbert_model/model.tflite"
-convert_to_tflite_model(nlu.model.model, save_file_path, conversion_mode="fp16_quantization")
+save_file_path = model_path + "/model.tflite"
+convert_to_tflite_model(nlu.model.model, save_file_path, conversion_mode="hybrid_quantization")
+print("Done")

examples/evaluate_tflite_transformer_nlu.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+"""
+@author: mwahdan
+"""
+
+# diasable the GPU
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+
+from dialognlu import TransformerNLU
+from dialognlu.readers.goo_format_reader import Reader
+import time
+
+
+num_process = 2
+
+
+model_path = "../saved_models/joint_distilbert_model"
+# model_path = "../saved_models/joint_trans_bert_model"
+# model_path = "../saved_models/joint_trans_albert_model"
+# model_path = "../saved_models/joint_trans_roberta_model"
+
+print("Loading model ...")
+nlu = TransformerNLU.load(model_path, quantized=True, num_process=num_process)
+
+print("Loading dataset ...")
+test_path = "../data/snips/test"
+test_dataset = Reader.read(test_path)
+
+print("Evaluating model ...")
+t1 = time.time()
+token_f1_score, tag_f1_score, report, acc = nlu.evaluate(test_dataset)
+t2 = time.time()
+
+print('Slot Classification Report:', report)
+print('Slot token f1_score = %f' % token_f1_score)
+print('Slot tag f1_score = %f' % tag_f1_score)
+print('Intent accuracy = %f' % acc)
+
+print("Using %d processes took %f seconds" % (num_process, t2 - t1))

examples/predict_tflite_nlu.py

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+"""
+@author: mwahdan
+"""
+
+# diasable the GPU
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+
+from dialognlu import TransformerNLU
+
+
+# model_path = "../saved_models/joint_distilbert_model"
+# model_path = "../saved_models/joint_trans_bert_model"
+# model_path = "../saved_models/joint_trans_albert_model"
+model_path = "../saved_models/joint_trans_roberta_model"
+
+print("Loading model ...")
+nlu = TransformerNLU.load(model_path, quantized=True, num_process=1)
+
+print("Prediction ...")
+utterance = "add sabrina salerno to the grime instrumentals playlist"
+print ("utterance: {}".format(utterance))
+result = nlu.predict(utterance)
+print ("result: {}".format(result))

examples/train_transformer_nlu.py

@@ -29,7 +29,7 @@
 pretrained_model_name_or_path = "distilbert-base-uncased"
 save_path = "../saved_models/joint_distilbert_model"
 
-epochs = 1 #3
+epochs = 3
 batch_size = 64
 
 
@@ -40,12 +40,14 @@
     "num_bert_fine_tune_layers": 10,
     "intent_loss_weight": 1.0,#0.2,
     "slots_loss_weight": 3.0,#2.0,
+
+    "max_length": 64, # You can set max_length (recommended) or leave it and it will be computed automatically based on longest training example
 }
 
 
 nlu = TransformerNLU.from_config(config)
 nlu.train(train_dataset, val_dataset, epochs, batch_size)
 
 print("Saving ...")
-nlu.save(save_path)
+nlu.save(save_path, save_tflite=True, conversion_mode="hybrid_quantization")
 print("Done")

setup.py

@@ -20,7 +20,7 @@
 
 setup(
     name="dialognlu",
-    version="0.1.0",
+    version="0.2.0",
     author="Mahmoud Wahdan",
     author_email="[email protected]",
     description="State-of-the-art Dialog NLU (Natural Language Understanding) Library with TensorFlow 2.x and keras",

src/dialognlu/__init__.py

@@ -1,5 +1,5 @@
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 from .nlu_components import TransformerNLU, BertNLU
 from .auto_nlu import AutoNLU

src/dialognlu/models/base_joint_trans.py

@@ -22,6 +22,7 @@ def __init__(self, config, trans_model=None, is_load=False):
         self.num_bert_fine_tune_layers = config.get('num_bert_fine_tune_layers', 10)
         self.intent_loss_weight = config.get('intent_loss_weight', 1.0)
         self.slots_loss_weight = config.get('slots_loss_weight', 3.0)
+        self.max_length = config.get('max_length')
 
         self.model_params = config
 
@@ -115,4 +116,65 @@ def load_model_by_class(klazz, load_folder_path, trans_model_name):
         new_model = klazz(model_params, trans_model=None, is_load=True)
         new_model.model = tf.keras.models.load_model(os.path.join(load_folder_path, trans_model_name))
         new_model.compile_model()
-        return new_model
+        return new_model
+
+
+class TfliteBaseJointTransformerModel:
+
+    def __init__(self, config):
+        self.config = config
+        self.slots_num = config['slots_num']
+        self.interpreter = None
+
+    def predict_slots_intent(self, x, slots_vectorizer, intent_vectorizer, remove_start_end=True,
+                             include_intent_prob=False):
+        # x = {k:v[0] for k,v in x.items()}
+        valid_positions = x["valid_positions"]
+        x["valid_positions"] = self.prepare_valid_positions(valid_positions)
+        y_slots, y_intent = self.predict(x)
+        slots = slots_vectorizer.inverse_transform(y_slots, valid_positions)
+        if remove_start_end:
+            slots = [x[1:-1] for x in slots]
+
+        if not include_intent_prob:
+            intents = np.array([intent_vectorizer.inverse_transform([np.argmax(i)])[0] for i in y_intent])
+        else:
+            intents = np.array([(intent_vectorizer.inverse_transform([np.argmax(i)])[0], round(float(np.max(i)), 4)) for i in y_intent])
+        return slots[0], intents[0]
+
+    def prepare_valid_positions(self, in_valid_positions):
+        in_valid_positions = np.expand_dims(in_valid_positions, axis=2)
+        in_valid_positions = np.tile(in_valid_positions, (1, 1, self.slots_num))
+        return in_valid_positions
+
+    def predict(self, inputs):
+        raise NotImplementedError()
+
+    @staticmethod
+    def load_model_by_class(clazz, path):
+        with open(os.path.join(path, 'params.json'), 'r') as json_file:
+            model_params = json.load(json_file)
+
+        new_model = clazz(model_params)
+        quant_model_file = os.path.join(path, 'model.tflite')
+        new_model.interpreter = tf.lite.Interpreter(model_path=str(quant_model_file), num_threads=1)
+        new_model.interpreter.allocate_tensors()
+        return new_model
+
+
+class TfliteBaseJointTransformer4inputsModel(TfliteBaseJointTransformerModel):
+
+    def __init__(self, config):
+        super(TfliteBaseJointTransformer4inputsModel, self).__init__(config)
+
+    def predict(self, inputs):
+        self.interpreter.set_tensor(self.interpreter.get_input_details()[0]["index"], inputs.get("input_word_ids").astype(np.int32))
+        self.interpreter.set_tensor(self.interpreter.get_input_details()[1]["index"], inputs.get("input_mask").astype(np.int32))
+        self.interpreter.set_tensor(self.interpreter.get_input_details()[2]["index"], inputs.get("input_type_ids").astype(np.int32))
+        self.interpreter.set_tensor(self.interpreter.get_input_details()[3]["index"], inputs.get("valid_positions").astype(np.float32))
+        output_index_0 = self.interpreter.get_output_details()[0]["index"]
+        output_index_1 = self.interpreter.get_output_details()[1]["index"]
+        self.interpreter.invoke()
+        intent = self.interpreter.get_tensor(output_index_0)
+        slots = self.interpreter.get_tensor(output_index_1)
+        return slots, intent

src/dialognlu/models/joint_trans_albert.py

@@ -6,7 +6,7 @@
 import tensorflow as tf
 from tensorflow.keras.models import Model
 from tensorflow.keras.layers import Input, Dense, Multiply, TimeDistributed
-from .base_joint_trans import BaseJointTransformerModel
+from .base_joint_trans import BaseJointTransformerModel, TfliteBaseJointTransformerModel, TfliteBaseJointTransformer4inputsModel
 
 
 class JointTransAlbertModel(BaseJointTransformerModel):
@@ -16,10 +16,10 @@ def __init__(self, config, trans_model=None, is_load=False):
 
 
     def build_model(self):
-        in_id = Input(shape=(None,), name='input_word_ids', dtype=tf.int32)
-        in_mask = Input(shape=(None,), name='input_mask', dtype=tf.int32)
-        in_segment = Input(shape=(None,), name='input_type_ids', dtype=tf.int32)
-        in_valid_positions = Input(shape=(None, self.slots_num), name='valid_positions')
+        in_id = Input(shape=(self.max_length), name='input_word_ids', dtype=tf.int32)
+        in_mask = Input(shape=(self.max_length), name='input_mask', dtype=tf.int32)
+        in_segment = Input(shape=(self.max_length), name='input_type_ids', dtype=tf.int32)
+        in_valid_positions = Input(shape=(self.max_length, self.slots_num), name='valid_positions')
         bert_inputs = [in_id, in_mask, in_segment]
         inputs = bert_inputs + [in_valid_positions]
 
@@ -41,3 +41,14 @@ def save(self, model_path):
     @staticmethod
     def load(load_folder_path):
         return BaseJointTransformerModel.load_model_by_class(JointTransAlbertModel, load_folder_path, 'joint_albert_model.h5')
+
+
+
+class TfliteJointTransAlbertModel(TfliteBaseJointTransformer4inputsModel):
+
+    def __init__(self, config):
+        super(TfliteJointTransAlbertModel, self).__init__(config)
+
+    @staticmethod
+    def load(path):
+        return TfliteBaseJointTransformerModel.load_model_by_class(TfliteJointTransAlbertModel, path)

src/dialognlu/models/joint_trans_bert.py

@@ -6,7 +6,8 @@
 import tensorflow as tf
 from tensorflow.keras.models import Model
 from tensorflow.keras.layers import Input, Dense, Multiply, TimeDistributed
-from .base_joint_trans import BaseJointTransformerModel
+from .base_joint_trans import BaseJointTransformerModel, TfliteBaseJointTransformerModel, TfliteBaseJointTransformer4inputsModel
+import numpy as np
 
 
 class JointTransBertModel(BaseJointTransformerModel):
@@ -16,10 +17,10 @@ def __init__(self, config, trans_model=None, is_load=False):
 
 
     def build_model(self):
-        in_id = Input(shape=(None,), name='input_word_ids', dtype=tf.int32)
-        in_mask = Input(shape=(None,), name='input_mask', dtype=tf.int32)
-        in_segment = Input(shape=(None,), name='input_type_ids', dtype=tf.int32)
-        in_valid_positions = Input(shape=(None, self.slots_num), name='valid_positions')
+        in_id = Input(shape=(self.max_length), name='input_word_ids', dtype=tf.int32)
+        in_mask = Input(shape=(self.max_length), name='input_mask', dtype=tf.int32)
+        in_segment = Input(shape=(self.max_length), name='input_type_ids', dtype=tf.int32)
+        in_valid_positions = Input(shape=(self.max_length, self.slots_num), name='valid_positions')
         bert_inputs = [in_id, in_mask, in_segment]
         inputs = bert_inputs + [in_valid_positions]
 
@@ -41,3 +42,14 @@ def save(self, model_path):
     @staticmethod    
     def load(load_folder_path):
         return BaseJointTransformerModel.load_model_by_class(JointTransBertModel, load_folder_path, 'joint_bert_model.h5')
+
+
+
+class TfliteJointTransBertModel(TfliteBaseJointTransformer4inputsModel):
+
+    def __init__(self, config):
+        super(TfliteJointTransBertModel, self).__init__(config)
+
+    @staticmethod
+    def load(path):
+        return TfliteBaseJointTransformerModel.load_model_by_class(TfliteJointTransBertModel, path)

src/dialognlu/models/joint_trans_distilbert.py

@@ -3,7 +3,7 @@
 @author: mwahdan
 """
 
-from .base_joint_trans import BaseJointTransformerModel
+from .base_joint_trans import BaseJointTransformerModel, TfliteBaseJointTransformerModel
 from .callbacks import F1Metrics
 import tensorflow as tf
 from tensorflow.keras.models import Model
@@ -18,9 +18,9 @@ def __init__(self, config, trans_model=None, is_load=False):
 
 
     def build_model(self):
-        in_id = Input(shape=(None,), name='input_word_ids', dtype=tf.int32)
-        in_mask = Input(shape=(None,), name='input_mask', dtype=tf.int32)
-        in_valid_positions = Input(shape=(None, self.slots_num), name='valid_positions')
+        in_id = Input(shape=(self.max_length), name='input_word_ids', dtype=tf.int32)
+        in_mask = Input(shape=(self.max_length), name='input_mask', dtype=tf.int32)
+        in_valid_positions = Input(shape=(self.max_length, self.slots_num), name='valid_positions')
         bert_inputs = [in_id, in_mask]
         inputs = bert_inputs + [in_valid_positions]
 
@@ -78,4 +78,26 @@ def save(self, model_path):
 
     @staticmethod
     def load(load_folder_path):
-        return BaseJointTransformerModel.load_model_by_class(JointTransDistilBertModel, load_folder_path, 'joint_distilbert_model.h5')
+        return BaseJointTransformerModel.load_model_by_class(JointTransDistilBertModel, load_folder_path, 'joint_distilbert_model.h5')
+
+
+
+class TfliteJointTransDistilBertModel(TfliteBaseJointTransformerModel):
+
+    def __init__(self, config):
+        super(TfliteJointTransDistilBertModel, self).__init__(config)
+
+    def predict(self, inputs):
+        self.interpreter.set_tensor(self.interpreter.get_input_details()[0]["index"], inputs.get("input_word_ids").astype(np.int32))
+        self.interpreter.set_tensor(self.interpreter.get_input_details()[1]["index"], inputs.get("input_mask").astype(np.int32))
+        self.interpreter.set_tensor(self.interpreter.get_input_details()[2]["index"], inputs.get("valid_positions").astype(np.float32))
+        output_index_0 = self.interpreter.get_output_details()[0]["index"]
+        output_index_1 = self.interpreter.get_output_details()[1]["index"]
+        self.interpreter.invoke()
+        intent = self.interpreter.get_tensor(output_index_0)
+        slots = self.interpreter.get_tensor(output_index_1)
+        return slots, intent
+
+    @staticmethod
+    def load(path):
+        return TfliteBaseJointTransformerModel.load_model_by_class(TfliteJointTransDistilBertModel, path)

src/dialognlu/models/joint_trans_roberta.py

@@ -3,10 +3,13 @@
 @author: mwahdan
 """
 
-from .base_joint_trans import BaseJointTransformerModel
 import tensorflow as tf
+from tensorflow.keras.layers import Dense, Input, Multiply, TimeDistributed
 from tensorflow.keras.models import Model
-from tensorflow.keras.layers import Input, Dense, Multiply, TimeDistributed
+
+from .base_joint_trans import (BaseJointTransformerModel,
+                               TfliteBaseJointTransformer4inputsModel,
+                               TfliteBaseJointTransformerModel)
 
 
 class JointTransRobertaModel(BaseJointTransformerModel):
@@ -16,10 +19,10 @@ def __init__(self, config, trans_model=None, is_load=False):
 
 
     def build_model(self):
-        in_id = Input(shape=(None,), name='input_word_ids', dtype=tf.int32)
-        in_mask = Input(shape=(None,), name='input_mask', dtype=tf.int32)
-        in_segment = Input(shape=(None,), name='input_type_ids', dtype=tf.int32)
-        in_valid_positions = Input(shape=(None, self.slots_num), name='valid_positions')
+        in_id = Input(shape=(self.max_length), name='input_word_ids', dtype=tf.int32)
+        in_mask = Input(shape=(self.max_length), name='input_mask', dtype=tf.int32)
+        in_segment = Input(shape=(self.max_length), name='input_type_ids', dtype=tf.int32)
+        in_valid_positions = Input(shape=(self.max_length, self.slots_num), name='valid_positions')
         bert_inputs = [in_id, in_mask, in_segment]
         inputs = bert_inputs + [in_valid_positions]
 
@@ -41,3 +44,13 @@ def save(self, model_path):
     @staticmethod
     def load(load_folder_path):
         return BaseJointTransformerModel.load_model_by_class(JointTransRobertaModel, load_folder_path, 'joint_roberta_model.h5')
+
+
+class TfliteJointTransRobertaModel(TfliteBaseJointTransformer4inputsModel):
+
+    def __init__(self, config):
+        super(TfliteJointTransRobertaModel, self).__init__(config)
+
+    @staticmethod
+    def load(path):
+        return TfliteBaseJointTransformerModel.load_model_by_class(TfliteJointTransRobertaModel, path)