Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: Aishwarya-Tonpe

superbench/benchmarks/micro_benchmarks/_export_torch_to_onnx.py

@@ -18,6 +18,7 @@
 
 class torch2onnxExporter():
     """PyTorch model to ONNX exporter."""
+
     def __init__(self):
         """Constructor."""
         from transformers import BertConfig, GPT2Config, LlamaConfig
@@ -314,95 +315,15 @@ def export_huggingface_model(self, model, model_name, batch_size=1, seq_length=5
             is_vision_model = main_input == 'pixel_values'
 
             if is_vision_model:
-                # Vision models: use pixel_values (batch_size, channels, height, width)
-                # Derive C/H/W from model config rather than hard-coding 3x224x224
-                num_channels = getattr(model.config, 'num_channels', 3)
-                image_size = getattr(model.config, 'image_size', 224)
-                if isinstance(image_size, (list, tuple)):
-                    img_h, img_w = image_size[0], image_size[1]
-                else:
-                    img_h, img_w = image_size, image_size
-
-                dummy_input = torch.randn(batch_size, num_channels, img_h, img_w, dtype=model_dtype, device=device)
-                input_names = ['pixel_values']
-                dynamic_axes = {'pixel_values': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
-
-                # Wrapper for vision models
-                class VisionModelWrapper(torch.nn.Module):
-                    def __init__(self, model):
-                        super().__init__()
-                        self.model = model
-
-                    def forward(self, pixel_values):
-                        outputs = self.model(pixel_values=pixel_values)
-                        if hasattr(outputs, 'logits'):
-                            return outputs.logits
-                        elif hasattr(outputs, 'last_hidden_state'):
-                            return outputs.last_hidden_state
-                        else:
-                            return outputs[0] if isinstance(outputs, (tuple, list)) else outputs
-
-                wrapped_model = VisionModelWrapper(model)
-                export_args = (dummy_input, )
+                wrapped_model, export_args, input_names, dynamic_axes = self._build_vision_export_inputs(
+                    model, batch_size, model_dtype, device
+                )
             else:
-                # NLP models: use input_ids and attention_mask
-                dummy_input = torch.ones((batch_size, seq_length), dtype=torch.int64, device=device)
-                attention_mask = torch.ones((batch_size, seq_length), dtype=torch.int64, device=device)
-                input_names = ['input_ids', 'attention_mask']
-                dynamic_axes = {
-                    'input_ids': {
-                        0: 'batch_size',
-                        1: 'seq_length'
-                    },
-                    'attention_mask': {
-                        0: 'batch_size',
-                        1: 'seq_length'
-                    },
-                    'output': {
-                        0: 'batch_size',
-                        1: 'seq_length'
-                    },
-                }
+                wrapped_model, export_args, input_names, dynamic_axes = self._build_nlp_export_inputs(
+                    model, batch_size, seq_length, device
+                )
 
-                # Wrapper for NLP models
-                class NLPModelWrapper(torch.nn.Module):
-                    def __init__(self, model):
-                        super().__init__()
-                        self.model = model
-
-                    def forward(self, input_ids, attention_mask):
-                        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
-                        if hasattr(outputs, 'logits'):
-                            return outputs.logits
-                        elif hasattr(outputs, 'last_hidden_state'):
-                            return outputs.last_hidden_state
-                        else:
-                            return outputs[0] if isinstance(outputs, (tuple, list)) else outputs
-
-                wrapped_model = NLPModelWrapper(model)
-                export_args = (dummy_input, attention_mask)
-
-            # Export to ONNX for large models (>2GB), use external data format
-            model_size_gb = sum(p.numel() * p.element_size() for p in model.parameters()) / (1024**3)
-            use_external_data = model_size_gb > 2.0
-
-            if use_external_data:
-                logger.info(f'Model size is {model_size_gb:.2f}GB, using external data format for ONNX export')
-
-            export_kwargs = {
-                'opset_version': 14,
-                'do_constant_folding': True,
-                'input_names': input_names,
-                'output_names': ['output'],
-                'dynamic_axes': dynamic_axes,
-            }
-            if use_external_data:
-                # PyTorch 2.8+ renamed 'use_external_data_format' to 'external_data'
-                sig = inspect.signature(torch.onnx.export)
-                if 'external_data' in sig.parameters:
-                    export_kwargs['external_data'] = True
-                else:
-                    export_kwargs['use_external_data_format'] = True
+            export_kwargs = self._build_onnx_export_kwargs(model, input_names, dynamic_axes)
 
             torch.onnx.export(
                 wrapped_model,
@@ -412,7 +333,7 @@ def forward(self, input_ids, attention_mask):
             )
 
             # Clean up
-            del dummy_input
+            del export_args
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
 
@@ -422,3 +343,97 @@ def forward(self, input_ids, attention_mask):
             logger.error(f'Failed to export HuggingFace model to ONNX: {str(e)}')
             logger.error(traceback.format_exc())
             return ''
+
+    def _build_vision_export_inputs(self, model, batch_size, model_dtype, device):
+        """Build the dummy inputs and wrapper module for exporting a vision HuggingFace model."""
+        # Vision models: use pixel_values (batch_size, channels, height, width)
+        # Derive C/H/W from model config rather than hard-coding 3x224x224
+        num_channels = getattr(model.config, 'num_channels', 3)
+        image_size = getattr(model.config, 'image_size', 224)
+        if isinstance(image_size, (list, tuple)):
+            img_h, img_w = image_size[0], image_size[1]
+        else:
+            img_h, img_w = image_size, image_size
+
+        dummy_input = torch.randn(batch_size, num_channels, img_h, img_w, dtype=model_dtype, device=device)
+        input_names = ['pixel_values']
+        dynamic_axes = {'pixel_values': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
+
+        class VisionModelWrapper(torch.nn.Module):
+
+            def __init__(self, model):
+                super().__init__()
+                self.model = model
+
+            def forward(self, pixel_values):
+                outputs = self.model(pixel_values=pixel_values)
+                if hasattr(outputs, 'logits'):
+                    return outputs.logits
+                elif hasattr(outputs, 'last_hidden_state'):
+                    return outputs.last_hidden_state
+                else:
+                    return outputs[0] if isinstance(outputs, (tuple, list)) else outputs
+
+        return VisionModelWrapper(model), (dummy_input, ), input_names, dynamic_axes
+
+    def _build_nlp_export_inputs(self, model, batch_size, seq_length, device):
+        """Build the dummy inputs and wrapper module for exporting an NLP HuggingFace model."""
+        # NLP models: use input_ids and attention_mask
+        dummy_input = torch.ones((batch_size, seq_length), dtype=torch.int64, device=device)
+        attention_mask = torch.ones((batch_size, seq_length), dtype=torch.int64, device=device)
+        input_names = ['input_ids', 'attention_mask']
+        dynamic_axes = {
+            'input_ids': {
+                0: 'batch_size',
+                1: 'seq_length'
+            },
+            'attention_mask': {
+                0: 'batch_size',
+                1: 'seq_length'
+            },
+            'output': {
+                0: 'batch_size',
+                1: 'seq_length'
+            },
+        }
+
+        class NLPModelWrapper(torch.nn.Module):
+
+            def __init__(self, model):
+                super().__init__()
+                self.model = model
+
+            def forward(self, input_ids, attention_mask):
+                outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+                if hasattr(outputs, 'logits'):
+                    return outputs.logits
+                elif hasattr(outputs, 'last_hidden_state'):
+                    return outputs.last_hidden_state
+                else:
+                    return outputs[0] if isinstance(outputs, (tuple, list)) else outputs
+
+        return NLPModelWrapper(model), (dummy_input, attention_mask), input_names, dynamic_axes
+
+    def _build_onnx_export_kwargs(self, model, input_names, dynamic_axes):
+        """Assemble torch.onnx.export kwargs, enabling external-data format for >2GB models."""
+        model_size_gb = sum(p.numel() * p.element_size() for p in model.parameters()) / (1024**3)
+        use_external_data = model_size_gb > 2.0
+
+        if use_external_data:
+            logger.info(f'Model size is {model_size_gb:.2f}GB, using external data format for ONNX export')
+
+        export_kwargs = {
+            'opset_version': 14,
+            'do_constant_folding': True,
+            'input_names': input_names,
+            'output_names': ['output'],
+            'dynamic_axes': dynamic_axes,
+        }
+        if use_external_data:
+            # PyTorch 2.8+ renamed 'use_external_data_format' to 'external_data'
+            sig = inspect.signature(torch.onnx.export)
+            if 'external_data' in sig.parameters:
+                export_kwargs['external_data'] = True
+            else:
+                export_kwargs['use_external_data_format'] = True
+        return export_kwargs

superbench/benchmarks/micro_benchmarks/huggingface_model_loader.py

@@ -85,6 +85,7 @@ class HuggingFaceModelLoader:
             ``False``; enabling this turns ``--model_identifier`` into an RCE
             sink, so it is opt-in only.
     """
+
     def __init__(
         self,
         cache_dir: Optional[str] = None,
@@ -150,19 +151,7 @@ def load_model(
         validate_model_identifier(model_identifier)
 
         try:
-            # Convert torch_dtype string to torch dtype
-            dtype = self._get_torch_dtype(torch_dtype) if torch_dtype else None
-
-            # Prepare loading kwargs
-            load_kwargs = {'cache_dir': self.cache_dir, 'revision': revision, **kwargs}
-
-            # Add token if available
-            if self.token:
-                load_kwargs['token'] = self.token
-
-            # Add dtype if specified
-            if dtype:
-                load_kwargs['torch_dtype'] = dtype
+            load_kwargs = self._build_load_kwargs(torch_dtype, revision, kwargs)
 
             # Load config (use pre-downloaded config if provided)
             if config is None:
@@ -173,35 +162,11 @@ def load_model(
             else:
                 logger.info('Using pre-downloaded model configuration.')
 
-            # Load tokenizer (may fail for some models, that's ok)
-            tokenizer = None
-            try:
-                logger.info('Loading tokenizer...')
-                tokenizer = AutoTokenizer.from_pretrained(
-                    model_identifier, trust_remote_code=self.allow_remote_code, **load_kwargs
-                )
-            except Exception as e:
-                logger.warning(f'Could not load tokenizer: {e}. Continuing without tokenizer.')
+            tokenizer = self._try_load_tokenizer(model_identifier, load_kwargs)
 
             # Load model
             logger.info(f'Loading model weights (dtype={torch_dtype}, device={device})...')
-            model_kwargs = load_kwargs.copy()
-            model_kwargs['trust_remote_code'] = self.allow_remote_code
-
-            # Handle device mapping for large models
-            effective_device_map = device_map
-            if device_map:
-                model_kwargs['device_map'] = device_map
-            elif device == 'cuda' and torch.cuda.is_available():
-                # Don't set device_map if device is explicitly cuda
-                pass
-            elif device != 'cpu':
-                model_kwargs['device_map'] = device
-                effective_device_map = device
-
-            # Pass pre-downloaded config to from_pretrained so any overrides take effect
-            if config is not None:
-                model_kwargs['config'] = config
+            model_kwargs, effective_device_map = self._build_model_kwargs(load_kwargs, device, device_map, config)
 
             try:
                 model = AutoModel.from_pretrained(model_identifier, **model_kwargs)
@@ -230,6 +195,52 @@ def load_model(
         except Exception as e:
             raise ModelLoadError(f"Unexpected error loading model '{model_identifier}': {e}") from e
 
+    def _build_load_kwargs(self, torch_dtype, revision, extra_kwargs):
+        """Assemble the base ``from_pretrained`` kwargs (cache_dir, token, dtype, revision)."""
+        dtype = self._get_torch_dtype(torch_dtype) if torch_dtype else None
+        load_kwargs = {'cache_dir': self.cache_dir, 'revision': revision, **extra_kwargs}
+        if self.token:
+            load_kwargs['token'] = self.token
+        if dtype:
+            load_kwargs['torch_dtype'] = dtype
+        return load_kwargs
+
+    def _try_load_tokenizer(self, model_identifier, load_kwargs):
+        """Attempt to load a tokenizer; return None if the model has no associated tokenizer."""
+        try:
+            logger.info('Loading tokenizer...')
+            return AutoTokenizer.from_pretrained(
+                model_identifier, trust_remote_code=self.allow_remote_code, **load_kwargs
+            )
+        except Exception as e:
+            logger.warning(f'Could not load tokenizer: {e}. Continuing without tokenizer.')
+            return None
+
+    def _build_model_kwargs(self, load_kwargs, device, device_map, config):
+        """Build model-loading kwargs and resolve the effective device_map.
+
+        Returns:
+            Tuple[dict, Optional[str]]: ``(model_kwargs, effective_device_map)``.
+        """
+        model_kwargs = load_kwargs.copy()
+        model_kwargs['trust_remote_code'] = self.allow_remote_code
+
+        effective_device_map = device_map
+        if device_map:
+            model_kwargs['device_map'] = device_map
+        elif device == 'cuda' and torch.cuda.is_available():
+            # Don't set device_map if device is explicitly cuda
+            pass
+        elif device != 'cpu':
+            model_kwargs['device_map'] = device
+            effective_device_map = device
+
+        # Pass pre-downloaded config to from_pretrained so any overrides take effect
+        if config is not None:
+            model_kwargs['config'] = config
+
+        return model_kwargs, effective_device_map
+
     def load_model_from_config(
         self,
         config: ModelSourceConfig,

superbench/benchmarks/micro_benchmarks/ort_inference_performance.py

@@ -232,7 +232,7 @@ def _preprocess_huggingface_models(self):
             self._hf_config = hf_config
 
             precision_str = self._args.precision.value if self._args.precision != Precision.INT8 else 'float32'
-            fits, param_m, est_gb, avail_gb = HuggingFaceModelLoader.check_memory_fits(
+            fits, _, _, _ = HuggingFaceModelLoader.check_memory_fits(
                 self._args.model_identifier, hf_config, precision_str, mode='inference', token=hf_token
             )
             if not fits:

tests/benchmarks/micro_benchmarks/test_huggingface_e2e.py

@@ -16,8 +16,10 @@
 
 pytest.importorskip('transformers')
 
-from superbench.benchmarks.micro_benchmarks.huggingface_model_loader import HuggingFaceModelLoader
-from superbench.benchmarks.micro_benchmarks.model_source_config import ModelSourceConfig
+# Imports below this point depend on `transformers` being available, so they
+# must be deferred until after the `importorskip` call above.
+from superbench.benchmarks.micro_benchmarks.huggingface_model_loader import HuggingFaceModelLoader    # noqa: E402
+from superbench.benchmarks.micro_benchmarks.model_source_config import ModelSourceConfig    # noqa: E402
 
 
 @pytest.mark.skipif(
@@ -26,6 +28,7 @@
 )
 class TestHuggingFaceE2E:
     """End-to-end tests for HuggingFace model loading."""
+
     @pytest.fixture
     def loader(self, tmp_path):
         """Create a loader instance with an isolated per-test cache dir."""