Parent class hf models (#1467)

Co-authored-by: Daniel King <[email protected]>
mosaicml · Aug 21, 2024 · c56401f · c56401f
1 parent 3648cef
commit c56401f
Show file tree

Hide file tree

Showing 5 changed files with 498 additions and 377 deletions.
diff --git a/llmfoundry/models/hf/__init__.py b/llmfoundry/models/hf/__init__.py
@@ -1,6 +1,7 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
+from llmfoundry.models.hf.hf_base import BaseHuggingFaceModel
 from llmfoundry.models.hf.hf_causal_lm import ComposerHFCausalLM
 from llmfoundry.models.hf.hf_fsdp import (
     prepare_hf_causal_lm_model_for_fsdp,
@@ -11,6 +12,7 @@
 from llmfoundry.models.hf.model_wrapper import HuggingFaceModelWithFSDP
 
 __all__ = [
+    'BaseHuggingFaceModel',
     'ComposerHFCausalLM',
     'ComposerHFT5',
     'prepare_hf_causal_lm_model_for_fsdp',