def save()

in src/datasets/utils/_dill.py [0:0]

• 30 lines of code
• 15 McCabe index (conditional complexity)

    def save(self, obj, save_persistent_id=True):
        obj_type = type(obj)
        if obj_type not in self.dispatch:
            if "regex" in sys.modules:
                import regex  # type: ignore

                if obj_type is regex.Pattern:
                    pklregister(obj_type)(_save_regexPattern)
            if "spacy" in sys.modules:
                import spacy  # type: ignore

                if issubclass(obj_type, spacy.Language):
                    pklregister(obj_type)(_save_spacyLanguage)
            if "tiktoken" in sys.modules:
                import tiktoken  # type: ignore

                if obj_type is tiktoken.Encoding:
                    pklregister(obj_type)(_save_tiktokenEncoding)
            if "torch" in sys.modules:
                import torch  # type: ignore

                if issubclass(obj_type, torch.Tensor):
                    pklregister(obj_type)(_save_torchTensor)

                if obj_type is torch.Generator:
                    pklregister(obj_type)(_save_torchGenerator)

                # Unwrap `torch.compile`-ed modules
                if issubclass(obj_type, torch.nn.Module):
                    obj = getattr(obj, "_orig_mod", obj)
            if "transformers" in sys.modules:
                import transformers  # type: ignore

                if issubclass(obj_type, transformers.PreTrainedTokenizerBase):
                    pklregister(obj_type)(_save_transformersPreTrainedTokenizerBase)

        # Unwrap `torch.compile`-ed functions
        if obj_type is FunctionType:
            obj = getattr(obj, "_torchdynamo_orig_callable", obj)
        dill.Pickler.save(self, obj, save_persistent_id=save_persistent_id)