precommit

mosaicml · Feb 2, 2024 · df17f5b · df17f5b
1 parent c8d1f6b
commit df17f5b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/tests/a_scripts/inference/test_convert_composer_to_hf.py b/tests/a_scripts/inference/test_convert_composer_to_hf.py
@@ -649,17 +649,19 @@ def test_huggingface_conversion_callback(
                 checkpoint_path = os.path.join(tmp_path, 'checkpoints',
                                                'huggingface',
                                                f'ba{batches_per_epoch}')
-                
+
                 if peft_config is not None:
-                    with open(os.path.join(checkpoint_path,
-                                        'adapter_config.json')) as _f:
+                    with open(
+                            os.path.join(checkpoint_path,
+                                         'adapter_config.json')) as _f:
                         adapter_config = json.load(_f)
 
-                    adapter_config['base_model_name_or_path'] = str(tmp_path /
-                                                                    'base-model')
+                    adapter_config['base_model_name_or_path'] = str(
+                        tmp_path / 'base-model')
 
-                    with open(os.path.join(checkpoint_path, 'adapter_config.json'),
-                            'w') as _f:
+                    with open(
+                            os.path.join(checkpoint_path,
+                                         'adapter_config.json'), 'w') as _f:
                         json.dump(adapter_config, _f)
 
                 # Load the last huggingface checkpoint

diff --git a/tests/models/hf/test_hf_peft_wrapping.py b/tests/models/hf/test_hf_peft_wrapping.py
@@ -9,7 +9,6 @@
 import torch
 import transformers
 from composer import Trainer
-from composer.models.huggingface import maybe_get_underlying_model
 from omegaconf import OmegaConf as om
 from peft import LoraConfig, get_peft_model