Formatting and remove logs

Author: jackzhxng

optimum/executorch/modeling.py

@@ -104,19 +104,19 @@ def __init__(self, models: Dict[str, "ExecuTorchModule"], config: "PretrainedCon
                 setattr(self, key, value)
 
         self.stats = Stats()
-        
+
         # Initialize cleanup tracking
         self._temp_dir = None
 
     def __del__(self):
         """Clean up temporary files when the model instance is destroyed."""
         self._cleanup_temp_resources()
-    
+
     def _cleanup_temp_resources(self):
         """Clean up temporary directory and files."""
-        if hasattr(self, '_temp_dir') and self._temp_dir is not None:
+        if hasattr(self, "_temp_dir") and self._temp_dir is not None:
             try:
-                if hasattr(self._temp_dir, 'cleanup'):
+                if hasattr(self._temp_dir, "cleanup"):
                     # It's a TemporaryDirectory object
                     logging.info(f"Cleaning up temporary directory: {self._temp_dir.name}")
                     self._temp_dir.cleanup()
@@ -216,18 +216,6 @@ def _from_pretrained(
             subfolder=subfolder,
             local_files_only=local_files_only,
         )
-
-        from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib  # noqa
-        from executorch.kernels import quantized  # noqa
-        from executorch.extension.pybindings.portable_lib import _get_operator_names
-        print("----------- LOADED OPS ----------")
-        print('\n'.join(_get_operator_names()))
-        print("---------------------------------")
-        model = _load_for_executorch(model_cache_path)
-        logging.info(
-            f"Loaded model from {model_cache_path} ({os.path.getsize(model_cache_path) / (1024 * 1024):.2f} MB)"
-        )
-
         return {default_file_name.removesuffix(_PTE_SUFFIX): model}
 
     @staticmethod
@@ -304,15 +292,6 @@ def _export(
         for name, _ in executorch_progs.items():
             models.update(cls._from_pretrained(save_dir_path, file_name=f"{name}.pte", config=config))
 
-        # Log temp directory info for debugging
-        logging.info(f"Created temporary directory: {save_dir_path}")
-        for name in executorch_progs.keys():
-            pte_file = save_dir_path / f"{name}.pte"
-            if pte_file.exists():
-                logging.info(f"PTE file exists at export: {pte_file} (size: {pte_file.stat().st_size} bytes)")
-            else:
-                logging.warning(f"PTE file missing at export: {pte_file}")
-        
         return models, save_dir
 
     def _save_pretrained(self, save_directory):
@@ -399,7 +378,9 @@ def from_pretrained(
                 **kwargs,
             )
         else:
-            logging.info(f"Pre-exported `.pte` artifact already exists in HuggingFace repo or provided file path for {model_id}, skipping export.")
+            logging.info(
+                f"Pre-exported `.pte` artifact already exists in HuggingFace repo or provided file path for {model_id}, skipping export."
+            )
             models_dict = {}
             for pte_file in pte_files:
                 models_dict.update(
@@ -418,12 +399,12 @@ def from_pretrained(
                 )
 
         model_instance = cls(models_dict, config)
-        
+
         # Store the TemporaryDirectory reference to prevent GC
         if temp_dir is not None:
             model_instance._temp_dir = temp_dir
             logging.info(f"Stored temp directory reference in model: {temp_dir.name}")
-            
+
         return model_instance
 
 
@@ -695,32 +676,12 @@ def forward(
         Returns:
             torch.Tensor: Logits output from the model.
         """
-        # Check if temp directory and PTE file still exist before forward pass
-        if hasattr(self, '_temp_dir') and self._temp_dir is not None:
-            temp_path = Path(self._temp_dir.name)
-            logging.info(f"Forward pass - temp directory exists: {temp_path.exists()}")
-            if temp_path.exists():
-                pte_files = list(temp_path.glob("*.pte"))
-                logging.info(f"Forward pass - PTE files found: {len(pte_files)}")
-                for pte_file in pte_files:
-                    logging.info(f"Forward pass - PTE file: {pte_file} exists: {pte_file.exists()}, size: {pte_file.stat().st_size if pte_file.exists() else 'N/A'}")
-            else:
-                logging.error(f"Forward pass - temp directory missing: {temp_path}")
-        else:
-            logging.info("Forward pass - no temp directory reference stored")
-
         self.stats.on_model_execution_start()
 
         try:
-            logging.info("Running forward()...")
             logits = self.model.forward((input_ids, cache_position))[0]
-            logging.info(f"logits from forward(): {logits}")
         except Exception as e:
             shapes = {name: val.shape for name, val in locals().items() if hasattr(val, "shape")}
-            logging.error(f"Forward pass failed - temp dir exists: {hasattr(self, '_temp_dir') and self._temp_dir is not None}")
-            if hasattr(self, '_temp_dir') and self._temp_dir is not None:
-                temp_path = Path(self._temp_dir.name)
-                logging.error(f"Forward pass failed - temp directory: {temp_path} exists: {temp_path.exists()}")
             print(f"Exception: {e}.\n{self.model.method_meta('forward')}\narg shapes: {shapes}")
             raise