Support native bert models with transformers class

Author: pwolnows

tools/accuracy_checker/accuracy_checker/launcher/pytorch_launcher.py

@@ -28,7 +28,7 @@
 MODULE_REGEX = r'(?:\w+)(?:(?:.\w+)*)'
 DEVICE_REGEX = r'(?P<device>cpu$|cuda)?'
 CHECKPOINT_URL_REGEX = r'^https?://.*\.pth(\?.*)?(#.*)?$'
-
+SCALAR_INPUTS = ('input_ids', 'input_mask', 'segment_ids', 'attention_mask', 'token_type_ids')
 
 class PyTorchLauncher(Launcher):
     __provider__ = 'pytorch'
@@ -99,26 +99,29 @@ def __init__(self, config_entry: dict, *args, **kwargs):
         module_kwargs = config_entry.get("module_kwargs", {})
         self.device = self.get_value_from_config('device')
         self.cuda = 'cuda' in self.device
+
         checkpoint = config_entry.get('checkpoint')
+        if checkpoint is None:
+            checkpoint = config_entry.get('checkpoint_url')
+
+        python_path = config_entry.get("python_path")
+
         if self.tranformers_class:
             self.module = self.load_tranformers_module(
-                config_entry['module']
+                config_entry['module'], python_path
             )
         else:
-            if checkpoint is None:
-                checkpoint = config_entry.get('checkpoint_url')
 
             self.module = self.load_module(
                 config_entry['module'],
                 module_args,
                 module_kwargs,
                 checkpoint,
                 config_entry.get('state_key'),
-                config_entry.get("python_path"),
+                python_path,
                 config_entry.get("init_method")
             )
 
-
         self._batch = self.get_value_from_config('batch')
         # torch modules does not have input information
         self._generate_inputs()
@@ -175,14 +178,15 @@ def load_module(self, model_cls, module_args, module_kwargs, checkpoint=None, st
 
             return self.prepare_module(module, model_cls)
 
-    def load_tranformers_module(self, pretrained_name):
-        import transformers # pylint: disable=C0415
-        model_class = getattr(transformers, self.tranformers_class)
-        module = model_class.from_pretrained(pretrained_name)
+    def load_tranformers_module(self, pretrained_name, python_path):
+        with append_to_path(python_path):
+            import transformers # pylint: disable=C0415
+            model_class = getattr(transformers, self.tranformers_class)
+            pretrained_model = python_path if python_path else pretrained_name
+            module = model_class.from_pretrained(pretrained_model)
 
         return self.prepare_module(module, model_class)
 
-
     def prepare_module(self, module, model_class):
         module.to('cuda' if self.cuda else 'cpu')
         module.eval()
@@ -191,16 +195,23 @@ def prepare_module(self, module, model_class):
             if hasattr(model_class, 'compile'):
                 module.compile()
             module = self._torch.compile(module, **self.compile_kwargs)
+
         return module
 
 
     def _convert_to_tensor(self, value, precision):
         if isinstance(value, self._torch.Tensor):
             return value
-        return self._torch.from_numpy(value.astype(np.float32 if not precision else precision)).to(self.device)
+        if precision is None:
+            precision =  np.float32
+
+        return self._torch.from_numpy(value.astype(precision)).to(self.device)
 
     def fit_to_input(self, data, layer_name, layout, precision, template=None):
 
+        if precision is None and layer_name in SCALAR_INPUTS:
+            precision = np.int64
+
         if layer_name == 'input' and isinstance(data[0], dict):
             tensor_dict = {}
             for key, val in data[0].items():
@@ -214,12 +225,14 @@ def fit_to_input(self, data, layer_name, layout, precision, template=None):
 
             return tensor_dict
 
-        if layout is not None:
+        data_shape = np.shape(data)
+
+        if layout is not None and len(data_shape) == len(layout):
             data = np.transpose(data, layout)
+        else:
+            data = np.array(data)
 
-        tensor = self._torch.from_numpy(data.astype(np.float32 if not precision else precision))
-        tensor = tensor.to(self.device)
-        return tensor
+        return self._convert_to_tensor(data, precision)
 
     def _convert_to_numpy(self, input_dict):
         numpy_dict = {}
@@ -230,25 +243,28 @@ def _convert_to_numpy(self, input_dict):
                 numpy_dict[key] = value
         return numpy_dict
 
+
+    def forward(self, outputs):
+        if hasattr(outputs, 'logits') and 'logits' in self.output_names:
+            return {'logits': outputs.logits}
+        elif hasattr(outputs, 'last_hidden_state') and 'last_hidden_state' in self.output_names:
+            return {'last_hidden_state': outputs.last_hidden_state}
+        return list(outputs)
+
     def predict(self, inputs, metadata=None, **kwargs):
         results = []
         with self._torch.no_grad():
             for batch_input in inputs:
-                if metadata[0].get('input_is_dict_type'):
+                if metadata[0].get('input_is_dict_type') or \
+                    (isinstance(batch_input, dict) and 'input' in batch_input):
                     outputs = self.module(batch_input['input'])
                 else:
-                    output = self.module(*batch_input.values())
-
-                    if 'logits' in self.output_names:
-                        result_dict = { 'logits': output.logits.detach().cpu().numpy() }
-                        results.append(result_dict)
-                        continue
-                    outputs = list(output)
+                    outputs = self.module(**batch_input)
 
                     for meta_ in metadata:
                         meta_['input_shape'] = {key: list(data.shape) for key, data in batch_input.items()}
 
-                if metadata[0].get('output_is_dict_type'):
+                if metadata[0].get('output_is_dict_type') or isinstance(outputs, dict):
                     result_dict = self._convert_to_numpy(outputs)
                 else:
                     result_dict = {