diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 2ec7745154..3e54ef2cd4 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -32,7 +32,7 @@ pass from huggingface_hub import HfFileSystem # https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading! -from unsloth_zoo.utils import Version +from unsloth_zoo.utils import Version, _get_dtype transformers_version = Version(transformers_version) SUPPORTS_FOURBIT = transformers_version >= Version("4.37") SUPPORTS_GEMMA = transformers_version >= Version("4.38") @@ -47,23 +47,6 @@ if SUPPORTS_GEMMA2: pass import torch -def _get_dtype(dtype): - __DTYPE_MAP = { - "float32": torch.float32, - torch.float32: torch.float32, - "float16": torch.float16, - torch.float16: torch.float16, - "bfloat16": torch.bfloat16, - torch.bfloat16: torch.bfloat16, - } - if dtype is None or dtype == None: return None - elif dtype in __DTYPE_MAP: return __DTYPE_MAP[dtype] - else: - print(f"Unsloth: {dtype} is not recognized, so we'll default to None") - return None - pass -pass - class FastLanguageModel(FastLlamaModel): @staticmethod