Upload blender_model.py

blender_model.py

@@ -0,0 +1,207 @@
+from transformers import (
+    AutoConfig,
+    BlenderbotSmallForConditionalGeneration,
+    logging
+)
+from transformers.modeling_outputs import (
+    Seq2SeqLMOutput,
+    BaseModelOutput,
+)
+from huggingface_hub import hf_hub_url, cached_download
+from onnxruntime import (GraphOptimizationLevel,
+                        InferenceSession,
+                        SessionOptions)
+
+from torch import from_numpy
+from torch.nn import Module
+from functools import reduce
+from operator import iconcat
+
+#supress huggingface warnings
+logging.set_verbosity_error()
+
+model_vocab_size=30000
+model_card="remzicam/xs_blenderbot_onnx"
+model_file_names=["blenderbot_small-90M-encoder-quantized.onnx",
+                "blenderbot_small-90M-decoder-quantized.onnx",
+                "blenderbot_small-90M-init-decoder-quantized.onnx"]
+
+class BlenderEncoder(Module):
+    def __init__(self, encoder_sess):
+        super().__init__()
+        self.encoder = encoder_sess
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask,
+        inputs_embeds=None,
+        head_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+
+        encoder_hidden_state = from_numpy(
+            self.encoder.run(
+                None,
+                {
+                    "input_ids": input_ids.cpu().numpy(),
+                    "attention_mask": attention_mask.cpu().numpy(),
+                },
+            )[0]
+        )
+
+        return BaseModelOutput(encoder_hidden_state)
+
+
+class BlenderDecoderInit(Module):
+    def __init__(self, decoder_sess):
+        super().__init__()
+        self.decoder = decoder_sess
+
+    def forward(self, input_ids, encoder_attention_mask, encoder_hidden_states):
+
+        decoder_outputs = self.decoder.run(
+            None,
+            {
+                "input_ids": input_ids.cpu().numpy(),
+                "encoder_attention_mask": encoder_attention_mask.cpu().numpy(),
+                "encoder_hidden_states": encoder_hidden_states.cpu().numpy(),
+            },
+        )
+
+        list_pkv = tuple(from_numpy(x) for x in decoder_outputs[1:])
+
+        out_past_key_values = tuple(
+            list_pkv[i : i + 4] for i in range(0, len(list_pkv), 4)
+        )
+
+        return from_numpy(decoder_outputs[0]), out_past_key_values
+
+
+class BlenderDecoder(Module):
+    def __init__(self, decoder_sess):
+        super().__init__()
+        self.decoder = decoder_sess
+
+    def forward(self, input_ids, attention_mask, encoder_output, past_key_values):
+
+        decoder_inputs = {
+            "input_ids": input_ids.cpu().numpy(),
+            "encoder_attention_mask": attention_mask.cpu().numpy(),
+        }
+
+        flat_past_key_values = reduce(iconcat, past_key_values, [])
+
+        past_key_values = {
+            f"pkv_{i}": pkv.cpu().numpy() for i, pkv in enumerate(flat_past_key_values)
+        }
+
+        decoder_outputs = self.decoder.run(None, {**decoder_inputs, **past_key_values})
+        # converts each value of the list to tensor from numpy
+        list_pkv = tuple(from_numpy(x) for x in decoder_outputs[1:])
+
+        # creates a tuple of tuples of shape 6x4 from the above tuple
+        out_past_key_values = tuple(
+            list_pkv[i : i + 4] for i in range(0, len(list_pkv), 4)
+        )
+
+        return from_numpy(decoder_outputs[0]), out_past_key_values
+
+
+class OnnxBlender(BlenderbotSmallForConditionalGeneration):
+    """creates a Blender model using onnx sessions (encode, decoder & init_decoder)"""
+
+    def __init__(self, onnx_model_sessions):
+        config = AutoConfig.from_pretrained("facebook/blenderbot_small-90M")
+        config.vocab_size=model_vocab_size
+        super().__init__(config)
+
+        assert len(onnx_model_sessions) == 3, "all three models should be given"
+
+        encoder_sess, decoder_sess, decoder_sess_init = onnx_model_sessions
+
+        self.encoder = BlenderEncoder(encoder_sess)
+        self.decoder = BlenderDecoder(decoder_sess)
+        self.decoder_init = BlenderDecoderInit(decoder_sess_init)
+
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+    
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        decoder_input_ids=None,
+        decoder_attention_mask=None,
+        head_mask=None,
+        decoder_head_mask=None,
+        cross_attn_head_mask=None,
+        encoder_outputs=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        decoder_inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        
+        encoder_hidden_states = encoder_outputs[0]
+       
+        if past_key_values is not None:
+            if decoder_input_ids is not None:
+                decoder_input_ids = decoder_input_ids[:, -1:]
+            if decoder_inputs_embeds is not None:
+                decoder_inputs_embeds = decoder_inputs_embeds[:, -1:]
+
+        if past_key_values is None:
+
+            # runs only for the first time:
+            init_onnx_outputs = self.decoder_init(
+                decoder_input_ids, attention_mask, encoder_hidden_states
+            )
+
+            logits, past_key_values = init_onnx_outputs
+
+        else:
+
+            onnx_outputs = self.decoder(
+                decoder_input_ids,
+                attention_mask,
+                encoder_hidden_states,
+                past_key_values,
+            )
+
+            logits, past_key_values = onnx_outputs
+
+        return Seq2SeqLMOutput(logits=logits, past_key_values=past_key_values)
+
+class ModelLoad:
+    def __init__(self, model_card,file_names):
+        self.model_card=model_card
+        self.file_names=file_names
+
+    def model_file_downloader(self,model_card,filename):
+        config_file_url = hf_hub_url(model_card, filename)
+        model_file = cached_download(config_file_url)
+        return model_file
+
+    def inference_session(self,file_name):
+        model_file=self.model_file_downloader(self.model_card,file_name)
+        options = SessionOptions()
+        options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
+        return InferenceSession(model_file,options=options)
+
+    def __call__(self,model_config):
+        model=model_config([*map(self.inference_session,
+                                self.file_names)])
+        return model
+
+model_loader=ModelLoad(model_card,model_file_names)
+blender_onnx_model=model_loader(OnnxBlender)