precommit

mosaicml · Sep 29, 2023 · 33f21a9 · 33f21a9
1 parent 5233e24
commit 33f21a9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -19,6 +19,7 @@
 from composer.utils import dist, format_name_with_dist_and_time, parse_uri
 from transformers import PreTrainedTokenizerBase, PreTrainedModel
 
+from llmfoundry.models.mpt import MPTConfig, MPTForCausalLM
 from llmfoundry.utils.huggingface_hub_utils import \
     edit_files_for_hf_compatibility
 

diff --git a/tests/test_hf_conversion_script.py b/tests/test_hf_conversion_script.py
@@ -12,8 +12,7 @@
 from composer.utils import dist, get_device
 
 from llmfoundry.callbacks import HuggingFaceCheckpointer
-from llmfoundry.models.mpt.modeling_mpt import (ComposerMPTCausalLM, MPTConfig,
-                                                MPTForCausalLM)
+from llmfoundry.models.mpt.modeling_mpt import ComposerMPTCausalLM
 
 # Add repo root to path so we can import scripts and test it
 repo_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))