Match transformers behavior with return_dict (#2269)

* fix

* fix

* more

Author: IlyasMoutawwakil

optimum/onnxruntime/modeling_diffusion.py

@@ -611,7 +611,7 @@ def forward(
         timestep_cond: Optional[Union[np.ndarray, torch.Tensor]] = None,
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         added_cond_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = False,
+        return_dict: bool = True,
     ):
         use_torch = isinstance(sample, torch.Tensor)
 
@@ -631,8 +631,8 @@ def forward(
         onnx_outputs = self.session.run(None, onnx_inputs)
         model_outputs = self.prepare_onnx_outputs(use_torch, *onnx_outputs)
 
-        if return_dict:
-            return model_outputs
+        if not return_dict:
+            return tuple(model_outputs.values())
 
         return ModelOutput(**model_outputs)
 
@@ -648,7 +648,7 @@ def forward(
         txt_ids: Optional[Union[np.ndarray, torch.Tensor]] = None,
         img_ids: Optional[Union[np.ndarray, torch.Tensor]] = None,
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = False,
+        return_dict: bool = True,
     ):
         use_torch = isinstance(hidden_states, torch.Tensor)
 
@@ -667,8 +667,8 @@ def forward(
         onnx_outputs = self.session.run(None, onnx_inputs)
         model_outputs = self.prepare_onnx_outputs(use_torch, *onnx_outputs)
 
-        if return_dict:
-            return model_outputs
+        if not return_dict:
+            return tuple(model_outputs.values())
 
         return ModelOutput(**model_outputs)
 
@@ -679,7 +679,7 @@ def forward(
         input_ids: Union[np.ndarray, torch.Tensor],
         attention_mask: Optional[Union[np.ndarray, torch.Tensor]] = None,
         output_hidden_states: Optional[bool] = None,
-        return_dict: bool = False,
+        return_dict: bool = True,
     ):
         use_torch = isinstance(input_ids, torch.Tensor)
 
@@ -700,8 +700,8 @@ def forward(
             for i in range(num_layers):
                 model_outputs.pop(f"hidden_states.{i}", None)
 
-        if return_dict:
-            return model_outputs
+        if not return_dict:
+            return tuple(model_outputs.values())
 
         return ModelOutput(**model_outputs)
 
@@ -722,7 +722,7 @@ def forward(
         self,
         sample: Union[np.ndarray, torch.Tensor],
         generator: Optional[torch.Generator] = None,
-        return_dict: bool = False,
+        return_dict: bool = True,
     ):
         use_torch = isinstance(sample, torch.Tensor)
 
@@ -740,8 +740,8 @@ def forward(
                 parameters=model_outputs.pop("latent_parameters")
             )
 
-        if return_dict:
-            return model_outputs
+        if not return_dict:
+            return tuple(model_outputs.values())
 
         return ModelOutput(**model_outputs)
 
@@ -762,7 +762,7 @@ def forward(
         self,
         latent_sample: Union[np.ndarray, torch.Tensor],
         generator: Optional[torch.Generator] = None,
-        return_dict: bool = False,
+        return_dict: bool = True,
     ):
         use_torch = isinstance(latent_sample, torch.Tensor)
 
@@ -775,8 +775,8 @@ def forward(
         if "latent_sample" in model_outputs:
             model_outputs["latents"] = model_outputs.pop("latent_sample")
 
-        if return_dict:
-            return model_outputs
+        if not return_dict:
+            return tuple(model_outputs.values())
 
         return ModelOutput(**model_outputs)
 

optimum/onnxruntime/modeling_ort.py

@@ -1077,7 +1077,7 @@ def forward(
         input_features: Optional[Union[torch.Tensor, np.ndarray]] = None,
         input_values: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -1134,8 +1134,8 @@ def forward(
                 # TODO: This allows to support sentence-transformers models (sentence embedding), but is not validated.
                 last_hidden_state = next(iter(model_outputs.values()))
 
-        if return_dict:
-            return {"last_hidden_state": last_hidden_state}
+        if not return_dict:
+            return (last_hidden_state,)
 
         # converts output to namedtuple for pipelines post-processing
         return BaseModelOutput(last_hidden_state=last_hidden_state)
@@ -1251,7 +1251,7 @@ def forward(
         attention_mask: Optional[Union[torch.Tensor, np.ndarray]] = None,
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -1288,8 +1288,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         # converts output to namedtuple for pipelines post-processing
         return MaskedLMOutput(logits=logits)
@@ -1353,7 +1353,7 @@ def forward(
         attention_mask: Optional[Union[torch.Tensor, np.ndarray]] = None,
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -1388,8 +1388,8 @@ def forward(
             start_logits = model_outputs["start_logits"]
             end_logits = model_outputs["end_logits"]
 
-        if return_dict:
-            return {"start_logits": start_logits, "end_logits": end_logits}
+        if not return_dict:
+            return (start_logits, end_logits)
 
         # converts output to namedtuple for pipelines post-processing
         return QuestionAnsweringModelOutput(start_logits=start_logits, end_logits=end_logits)
@@ -1468,7 +1468,7 @@ def forward(
         attention_mask: Optional[Union[torch.Tensor, np.ndarray]] = None,
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -1505,8 +1505,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         # converts output to namedtuple for pipelines post-processing
         return SequenceClassifierOutput(logits=logits)
@@ -1571,7 +1571,7 @@ def forward(
         attention_mask: Optional[Union[torch.Tensor, np.ndarray]] = None,
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -1608,8 +1608,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         return TokenClassifierOutput(logits=logits)
 
@@ -1667,7 +1667,7 @@ def forward(
         attention_mask: Optional[Union[torch.Tensor, np.ndarray]] = None,
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -1704,8 +1704,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         # converts output to namedtuple for pipelines post-processing
         return MultipleChoiceModelOutput(logits=logits)
@@ -1770,7 +1770,7 @@ def forward(
         self,
         pixel_values: Union[torch.Tensor, np.ndarray],
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -1802,8 +1802,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         # converts output to namedtuple for pipelines post-processing
         return ImageClassifierOutput(logits=logits)
@@ -1868,7 +1868,7 @@ def forward(
         self,
         pixel_values: Union[torch.Tensor, np.ndarray],
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -1900,8 +1900,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         # converts output to namedtuple for pipelines post-processing
         return SemanticSegmenterOutput(logits=logits)
@@ -1996,7 +1996,7 @@ def forward(
         attention_mask: Optional[Union[torch.Tensor, np.ndarray]] = None,
         input_features: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -2038,8 +2038,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         # converts output to namedtuple for pipelines post-processing
         return SequenceClassifierOutput(logits=logits)
@@ -2092,7 +2092,7 @@ def forward(
         self,
         input_values: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -2134,8 +2134,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         # converts output to namedtuple for pipelines post-processing
         return CausalLMOutput(logits=logits)
@@ -2196,7 +2196,7 @@ def forward(
         self,
         input_values: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -2231,8 +2231,8 @@ def forward(
             logits = model_outputs["logits"]
             embeddings = model_outputs["embeddings"]
 
-        if return_dict:
-            return {"logits": logits, "embeddings": embeddings}
+        if not return_dict:
+            return (logits, embeddings)
 
         # converts output to namedtuple for pipelines post-processing
         return XVectorOutput(logits=logits, embeddings=embeddings)
@@ -2285,7 +2285,7 @@ def forward(
         self,
         input_values: Optional[Union[torch.Tensor, np.ndarray]] = None,
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -2305,8 +2305,8 @@ def forward(
 
             logits = model_outputs["logits"]
 
-        if return_dict:
-            return {"logits": logits}
+        if not return_dict:
+            return (logits,)
 
         # converts output to namedtuple for pipelines post-processing
         return TokenClassifierOutput(logits=logits)
@@ -2353,7 +2353,7 @@ def forward(
         self,
         pixel_values: Union[torch.Tensor, np.ndarray],
         *,
-        return_dict: bool = False,
+        return_dict: bool = True,
         **kwargs,
     ):
         # Warn about any unexpected kwargs using the helper method
@@ -2390,8 +2390,8 @@ def forward(
             model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs)
             reconstruction = model_outputs["reconstruction"]
 
-        if return_dict:
-            return {"reconstruction": reconstruction}
+        if not return_dict:
+            return (reconstruction,)
 
         return ImageSuperResolutionOutput(reconstruction=reconstruction)
 

tests/onnxruntime/test_modeling.py

@@ -2138,17 +2138,16 @@ def test_compare_to_transformers(self, model_arch):
 
         for input_type in ["pt", "np"]:
             tokens = tokenizer(text, return_tensors=input_type)
-            # Test default behavior (return_dict=False)
+            # Test default behavior (return_dict=True)
             onnx_outputs = onnx_model(**tokens)
             self.assertIsInstance(onnx_outputs, BaseModelOutput)
             self.assertIn("last_hidden_state", onnx_outputs)
             self.assertIsInstance(onnx_outputs.last_hidden_state, self.TENSOR_ALIAS_TO_TYPE[input_type])
 
-            # Test return_dict=True
-            onnx_outputs_dict = onnx_model(**tokens, return_dict=True)
-            self.assertIsInstance(onnx_outputs_dict, dict)
-            self.assertIn("last_hidden_state", onnx_outputs_dict)
-            self.assertIsInstance(onnx_outputs_dict["last_hidden_state"], self.TENSOR_ALIAS_TO_TYPE[input_type])
+            # Test return_dict=False
+            onnx_outputs_dict = onnx_model(**tokens, return_dict=False)
+            self.assertIsInstance(onnx_outputs_dict, tuple)
+            self.assertIsInstance(onnx_outputs_dict[0], self.TENSOR_ALIAS_TO_TYPE[input_type])
 
             # compare tensor outputs
             torch.testing.assert_close(