From 0569cf98cd355a2155668b80cd24c7937462b250 Mon Sep 17 00:00:00 2001 From: Brian Lester Date: Mon, 13 May 2024 16:38:09 -0400 Subject: [PATCH] Implement the ability to infer the checkpoint type Currently this is just based on the file path. In order to make the sniffing extensible, it is implemented as a plugin. In order to do that, and to avoid having to import a deep learning framework to do sniffing, the sniffers are separated out into their own files. This lead to the question of why are checkpoint plugs stored in the main repo? If we move them to the actual plugins area we can continue to support them while not needing to install them if a framework isn't going to be used. Now when you want to git theta with a specific framework you can either run `pip install git-theta[framework]` or run `pip install git-theta-checkpoints-framework` for your framework. --- git_theta/checkpoints/__init__.py | 1 + git_theta/checkpoints/base.py | 18 +++++ git_theta/scripts/git_theta_filter.py | 4 ++ git_theta/utils.py | 2 +- .../git_theta_checkpoints_flax/__init__.py | 1 + .../git_theta_checkpoints_flax/checkpoints.py | 0 .../git_theta_checkpoints_flax/sniffer.py | 9 +++ plugins/checkpoints/flax/setup.py | 65 ++++++++++++++++++ .../git_theta_checkpoints_pytorch/__init__.py | 1 + .../checkpoints.py | 0 .../git_theta_checkpoints_pytorch/sniffer.py | 15 ++++ plugins/checkpoints/pytorch/setup.py | 64 +++++++++++++++++ .../__init__.py | 1 + .../checkpoints.py | 0 .../flycheck_checkpoints.py | 44 ++++++++++++ .../sniffer.py | 8 +++ plugins/checkpoints/safetensors/setup.py | 64 +++++++++++++++++ .../__init__.py | 1 + .../checkpoints.py | 0 .../sniffer.py | 14 ++++ plugins/checkpoints/tensorflow/setup.py | 68 +++++++++++++++++++ setup.py | 24 ++----- 22 files changed, 386 insertions(+), 18 deletions(-) create mode 100644 plugins/checkpoints/flax/git_theta_checkpoints_flax/__init__.py rename git_theta/checkpoints/flax_checkpoint.py => plugins/checkpoints/flax/git_theta_checkpoints_flax/checkpoints.py (100%) create mode 100644 plugins/checkpoints/flax/git_theta_checkpoints_flax/sniffer.py create mode 100644 plugins/checkpoints/flax/setup.py create mode 100644 plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/__init__.py rename git_theta/checkpoints/pickled_dict_checkpoint.py => plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/checkpoints.py (100%) create mode 100644 plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/sniffer.py create mode 100644 plugins/checkpoints/pytorch/setup.py create mode 100644 plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/__init__.py rename git_theta/checkpoints/safetensors_checkpoint.py => plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/checkpoints.py (100%) create mode 100644 plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/flycheck_checkpoints.py create mode 100644 plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/sniffer.py create mode 100644 plugins/checkpoints/safetensors/setup.py create mode 100644 plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/__init__.py rename git_theta/checkpoints/tensorflow_checkpoint.py => plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/checkpoints.py (100%) create mode 100644 plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/sniffer.py create mode 100644 plugins/checkpoints/tensorflow/setup.py diff --git a/git_theta/checkpoints/__init__.py b/git_theta/checkpoints/__init__.py index 731abcc..697255d 100644 --- a/git_theta/checkpoints/__init__.py +++ b/git_theta/checkpoints/__init__.py @@ -4,4 +4,5 @@ Checkpoint, get_checkpoint_handler, get_checkpoint_handler_name, + sniff_checkpoint, ) diff --git a/git_theta/checkpoints/base.py b/git_theta/checkpoints/base.py index c9ac39e..2f0303c 100644 --- a/git_theta/checkpoints/base.py +++ b/git_theta/checkpoints/base.py @@ -157,3 +157,21 @@ def get_checkpoint_handler(checkpoint_type: Optional[str] = None) -> Checkpoint: checkpoint_type = get_checkpoint_handler_name(checkpoint_type) discovered_plugins = entry_points(group="git_theta.plugins.checkpoints") return discovered_plugins[checkpoint_type].load() + + +def sniff_checkpoint(checkpoint_path) -> str: + """En""" + discovered_plugins = entry_points(group="git_theta.plugins.checkpoint.sniffers") + loaded_plugins = {ep.name: ep.load() for ep in discovered_plugins} + logger = logging.getLogger("git_theta") + logger.debug( + f"Sniffing {checkpoint_path} to infer which deep learning framework it is." + ) + for ckpt_type, ckpt_sniffer in loaded_plugins.items(): + logger.debug(f"Checking if {checkpoint_path} is a {ckpt_type} checkpoint.") + if ckpt_sniffer(checkpoint_path): + logger.debug( + f"Determined that {checkpoint_path} is a {ckpt_type} checkpoint." + ) + return ckpt_type + raise ValueError(f"Couldn't determine checkpoint type for {checkpoint_path}") diff --git a/git_theta/scripts/git_theta_filter.py b/git_theta/scripts/git_theta_filter.py index 63c666a..d18de9f 100755 --- a/git_theta/scripts/git_theta_filter.py +++ b/git_theta/scripts/git_theta_filter.py @@ -36,6 +36,8 @@ def run_clean(args): """ logger = logging.getLogger("git_theta") logger.debug(f"Running clean filter on {args.file}") + if EnvVarConstants.CHECKPOINT_TYPE == "sniff": + EnvVarConstants.CHECKPOINT_TYPE = checkpoints.sniff_checkpoints() repo = git_utils.get_git_repo() checkpoint_handler = checkpoints.get_checkpoint_handler() if EnvVarConstants.LOW_MEMORY: @@ -74,6 +76,8 @@ def run_smudge(args): """ logger = logging.getLogger("git_theta") logger.debug(f"Running smudge filter on {args.file}") + if EnvVarConstants.CHECKPOINT_TYPE == "sniff": + EnvVarConstants.CHECKPOINT_TYPE = checkpoints.sniff_checkpoints() repo = git_utils.get_git_repo() curr_metadata = metadata.Metadata.from_file(sys.stdin) diff --git a/git_theta/utils.py b/git_theta/utils.py index 6c6e921..cb1bf8f 100644 --- a/git_theta/utils.py +++ b/git_theta/utils.py @@ -71,7 +71,7 @@ def __get__(self, obj, objtype=None): class EnvVarConstants: - CHECKPOINT_TYPE = EnvVar(name="GIT_THETA_CHECKPOINT_TYPE", default="pytorch") + CHECKPOINT_TYPE = EnvVar(name="GIT_THETA_CHECKPOINT_TYPE", default="sniff") UPDATE_TYPE = EnvVar(name="GIT_THETA_UPDATE_TYPE", default="dense") UPDATE_DATA_PATH = EnvVar(name="GIT_THETA_UPDATE_DATA_PATH", default="") PARAMETER_ATOL = EnvVar(name="GIT_THETA_PARAMETER_ATOL", default=1e-8) diff --git a/plugins/checkpoints/flax/git_theta_checkpoints_flax/__init__.py b/plugins/checkpoints/flax/git_theta_checkpoints_flax/__init__.py new file mode 100644 index 0000000..d3ec452 --- /dev/null +++ b/plugins/checkpoints/flax/git_theta_checkpoints_flax/__init__.py @@ -0,0 +1 @@ +__version__ = "0.2.0" diff --git a/git_theta/checkpoints/flax_checkpoint.py b/plugins/checkpoints/flax/git_theta_checkpoints_flax/checkpoints.py similarity index 100% rename from git_theta/checkpoints/flax_checkpoint.py rename to plugins/checkpoints/flax/git_theta_checkpoints_flax/checkpoints.py diff --git a/plugins/checkpoints/flax/git_theta_checkpoints_flax/sniffer.py b/plugins/checkpoints/flax/git_theta_checkpoints_flax/sniffer.py new file mode 100644 index 0000000..95c741a --- /dev/null +++ b/plugins/checkpoints/flax/git_theta_checkpoints_flax/sniffer.py @@ -0,0 +1,9 @@ +"""Infer if a checkpoint is flax based. + +We put this in a different file to avoid importing dl frameworks for file sniffing. +""" + + +def flax_sniffer(checkpoint_path: str) -> bool: + # TODO: Check if the actual value is msgpack based on magic numbers? + return checkpoint_path.endswith(".flax") diff --git a/plugins/checkpoints/flax/setup.py b/plugins/checkpoints/flax/setup.py new file mode 100644 index 0000000..4ffcc1e --- /dev/null +++ b/plugins/checkpoints/flax/setup.py @@ -0,0 +1,65 @@ +"""Plugin to support the flax checkpoint format.""" + +import ast +import os + +from setuptools import setup + + +def get_version(file_name: str, version_variable: str = "__version__") -> str: + """Find the version by walking the AST to avoid duplication. + + Parameters + ---------- + file_name : str + The file we are parsing to get the version string from. + version_variable : str + The variable name that holds the version string. + + Raises + ------ + ValueError + If there was no assignment to version_variable in file_name. + + Returns + ------- + version_string : str + The version string parsed from file_name_name. + """ + with open(file_name) as f: + tree = ast.parse(f.read()) + # Look at all assignment nodes that happen in the ast. If the variable + # name matches the given parameter, grab the value (which will be + # the version string we are looking for). + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + if node.targets[0].id == version_variable: + return node.value.s + raise ValueError( + f"Could not find an assignment to {version_variable} " f"within '{file_name}'" + ) + + +setup( + name="git_theta_checkpoints_flax", + description="Plugin to support the flax checkpoint format.", + install_requires=[ + # "git_theta", + "flax", + "jax", + ], + version=get_version("git_theta_checkpoints_flax/__init__.py"), + packages=[ + "git_theta_checkpoints_flax", + ], + author="Brian Lester", + entry_points={ + "git_theta.plugins.checkpoints": [ + "flax = git_theta_checkpoints_flax.checkpoints:FlaxCheckpoint", + "flax-checkpoint = git_theta_checkpoints_flax.checkpoints:FlaxCheckpoint", + ], + "git_theta.plugins.checkpoint.sniffers": [ + "flax = git_theta_checkpoints_flax.sniffer:flax_sniffer", + ], + }, +) diff --git a/plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/__init__.py b/plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/__init__.py new file mode 100644 index 0000000..d3ec452 --- /dev/null +++ b/plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/__init__.py @@ -0,0 +1 @@ +__version__ = "0.2.0" diff --git a/git_theta/checkpoints/pickled_dict_checkpoint.py b/plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/checkpoints.py similarity index 100% rename from git_theta/checkpoints/pickled_dict_checkpoint.py rename to plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/checkpoints.py diff --git a/plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/sniffer.py b/plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/sniffer.py new file mode 100644 index 0000000..566d20e --- /dev/null +++ b/plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/sniffer.py @@ -0,0 +1,15 @@ +"""Infer if a checkpoint is pytorch based. + +We put this in a different file to avoid importing dl frameworks for file sniffing. +""" + +import re + + +def pytorch_sniffer(checkpoint_path: str) -> bool: + # Many checkpoints on HuggingFace Hub are named this. + if checkpoint_path == "pytorch_model.bin": + return True + if re.search(r"\.py?t$", checkpoint_path): + return True + return False diff --git a/plugins/checkpoints/pytorch/setup.py b/plugins/checkpoints/pytorch/setup.py new file mode 100644 index 0000000..b1d75e2 --- /dev/null +++ b/plugins/checkpoints/pytorch/setup.py @@ -0,0 +1,64 @@ +"""Plugin to support the safetensor format.""" + +import ast +import os + +from setuptools import setup + + +def get_version(file_name: str, version_variable: str = "__version__") -> str: + """Find the version by walking the AST to avoid duplication. + + Parameters + ---------- + file_name : str + The file we are parsing to get the version string from. + version_variable : str + The variable name that holds the version string. + + Raises + ------ + ValueError + If there was no assignment to version_variable in file_name. + + Returns + ------- + version_string : str + The version string parsed from file_name_name. + """ + with open(file_name) as f: + tree = ast.parse(f.read()) + # Look at all assignment nodes that happen in the ast. If the variable + # name matches the given parameter, grab the value (which will be + # the version string we are looking for). + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + if node.targets[0].id == version_variable: + return node.value.s + raise ValueError( + f"Could not find an assignment to {version_variable} " f"within '{file_name}'" + ) + + +setup( + name="git_theta_checkpoints_pytorch", + description="Plugin to support the pytorch checkpoint format.", + install_requires=[ + # "git_theta", + "torch", + ], + version=get_version("git_theta_checkpoints_pytorch/__init__.py"), + packages=[ + "git_theta_checkpoints_pytorch", + ], + author="Brian Lester", + entry_points={ + "git_theta.plugins.checkpoints": [ + "pytorch = git_theta_checkpoints_pytorch.checkpoints:PickledDictCheckpoint", + "pickled-dict = git_theta_checkpoints_pytorch.checkpoints:PickledDictCheckpoint", + ], + "git_theta.plugins.checkpoint.sniffers": [ + "pytorch = git_theta_checkpoints_pytorch.sniffer:pytorch_sniffer", + ], + }, +) diff --git a/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/__init__.py b/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/__init__.py new file mode 100644 index 0000000..d3ec452 --- /dev/null +++ b/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/__init__.py @@ -0,0 +1 @@ +__version__ = "0.2.0" diff --git a/git_theta/checkpoints/safetensors_checkpoint.py b/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/checkpoints.py similarity index 100% rename from git_theta/checkpoints/safetensors_checkpoint.py rename to plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/checkpoints.py diff --git a/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/flycheck_checkpoints.py b/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/flycheck_checkpoints.py new file mode 100644 index 0000000..543e627 --- /dev/null +++ b/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/flycheck_checkpoints.py @@ -0,0 +1,44 @@ +"""Checkpoint using the HF safetensors format. + +safetensors has the ability to write model checkpoint from "dl-native" -> "safetensors" +and read "safetensors" -> any "dl-native" framework, not just the one that was +used to write it. Therefore, we read/write with their numpy API. +""" + +import safetensors.numpy +from file_or_name import file_or_name + +from git_theta.checkpoints import Checkpoint + + +# TODO(bdlester): Can we leverage the lazying loading ability to make things faster? +class SafeTensorsCheckpoint(Checkpoint): + """Class for r/w of the safetensors format. https://github.com/huggingface/safetensors""" + + name: str = "safetensors" + + @classmethod + @file_or_name(checkpoint_path="rb") + def load(cls, checkpoint_path: str): + # Note that we use the numpy as the framework because we don't care what + # their downstream dl framework is, we only want the results back as + # numpy arrays. + return safetensors.numpy.load(checkpoint_path.read()) + + @file_or_name(checkpoint_path="wb") + def save(self, checkpoint_path: str): + # Note, git theta uses numpy internally, so we save using the numpy api, + # regardless of the original framework they used to write the checkpoint. + checkpoint_dict = self.to_framework() + checkpoint_path.write(safetensors.numpy.save(checkpoint_dict)) + + def to_framework(self): + return self + + @classmethod + def from_framework(cls, model_dict): + return cls(model_dict) + + +def safetensors_sniffer(checkpoint_path: str) -> bool: + return checkpoint_path.endswith(".safetensors") diff --git a/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/sniffer.py b/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/sniffer.py new file mode 100644 index 0000000..ee49a94 --- /dev/null +++ b/plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/sniffer.py @@ -0,0 +1,8 @@ +"""Infer if a checkpoint is safetensors based. + +We put this in a different file to avoid importing dl frameworks for file sniffing. +""" + + +def safetensors_sniffer(checkpoint_path: str) -> bool: + return checkpoint_path.endswith(".safetensors") diff --git a/plugins/checkpoints/safetensors/setup.py b/plugins/checkpoints/safetensors/setup.py new file mode 100644 index 0000000..a3ce5b2 --- /dev/null +++ b/plugins/checkpoints/safetensors/setup.py @@ -0,0 +1,64 @@ +"""Plugin to support the safetensor format.""" + +import ast +import os + +from setuptools import setup + + +def get_version(file_name: str, version_variable: str = "__version__") -> str: + """Find the version by walking the AST to avoid duplication. + + Parameters + ---------- + file_name : str + The file we are parsing to get the version string from. + version_variable : str + The variable name that holds the version string. + + Raises + ------ + ValueError + If there was no assignment to version_variable in file_name. + + Returns + ------- + version_string : str + The version string parsed from file_name_name. + """ + with open(file_name) as f: + tree = ast.parse(f.read()) + # Look at all assignment nodes that happen in the ast. If the variable + # name matches the given parameter, grab the value (which will be + # the version string we are looking for). + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + if node.targets[0].id == version_variable: + return node.value.s + raise ValueError( + f"Could not find an assignment to {version_variable} " f"within '{file_name}'" + ) + + +setup( + name="git_theta_checkpoints_safetensors", + description="Plugin to support the safetensors checkpoint format.", + install_requires=[ + # "git_theta", + "safetensors", + ], + version=get_version("git_theta_checkpoints_safetensors/__init__.py"), + packages=[ + "git_theta_checkpoints_safetensors", + ], + author="Brian Lester", + entry_points={ + "git_theta.plugins.checkpoints": [ + "safetensors = git_theta_checkpoints_safetensors.checkpoints:SafeTensorsCheckpoint", + "safetensors-checkpoint = git_theta_checkpoints_safetensors.checkpoints:SafeTensorsCheckpoint", + ], + "git_theta.plugins.checkpoint.sniffers": [ + "safetensors = git_theta_checkpoints_safetensors.sniffer:safetensors_sniffer", + ], + }, +) diff --git a/plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/__init__.py b/plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/__init__.py new file mode 100644 index 0000000..d3ec452 --- /dev/null +++ b/plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/__init__.py @@ -0,0 +1 @@ +__version__ = "0.2.0" diff --git a/git_theta/checkpoints/tensorflow_checkpoint.py b/plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/checkpoints.py similarity index 100% rename from git_theta/checkpoints/tensorflow_checkpoint.py rename to plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/checkpoints.py diff --git a/plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/sniffer.py b/plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/sniffer.py new file mode 100644 index 0000000..f10aab9 --- /dev/null +++ b/plugins/checkpoints/tensorflow/git_theta_checkpoints_tensorflow/sniffer.py @@ -0,0 +1,14 @@ +"""Infer if a checkpoint is tensorflow based. + +We put this in a different file to avoid importing dl frameworks for file sniffing. +""" + + +def tensorflow_sniffer(checkpoint_path: str) -> bool: + return checkpoint_path.endswith(".tf") + + +# TODO: Add support for detecting saved models. +def saved_model_sniffer(checkpoint_path: str) -> bool: + # We don't support saved models yet. + return False diff --git a/plugins/checkpoints/tensorflow/setup.py b/plugins/checkpoints/tensorflow/setup.py new file mode 100644 index 0000000..ed34fe9 --- /dev/null +++ b/plugins/checkpoints/tensorflow/setup.py @@ -0,0 +1,68 @@ +"""Plugin to support the tensorflow checkpoint format.""" + +import ast +import os + +from setuptools import setup + + +def get_version(file_name: str, version_variable: str = "__version__") -> str: + """Find the version by walking the AST to avoid duplication. + + Parameters + ---------- + file_name : str + The file we are parsing to get the version string from. + version_variable : str + The variable name that holds the version string. + + Raises + ------ + ValueError + If there was no assignment to version_variable in file_name. + + Returns + ------- + version_string : str + The version string parsed from file_name_name. + """ + with open(file_name) as f: + tree = ast.parse(f.read()) + # Look at all assignment nodes that happen in the ast. If the variable + # name matches the given parameter, grab the value (which will be + # the version string we are looking for). + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + if node.targets[0].id == version_variable: + return node.value.s + raise ValueError( + f"Could not find an assignment to {version_variable} " f"within '{file_name}'" + ) + + +setup( + name="git_theta_checkpoints_tensorflow", + description="Plugin to support the tensorflow checkpoint format.", + install_requires=[ + # "git_theta", + "tensorflow", + ], + version=get_version("git_theta_checkpoints_tensorflow/__init__.py"), + packages=[ + "git_theta_checkpoints_tensorflow", + ], + author="Brian Lester", + entry_points={ + "git_theta.plugins.checkpoints": [ + "tf = git_theta_checkpoints_tensorflow.checkpoints:TensorFlowCheckpoint", + "tensorflow = git_theta_checkpoints_tensorflow.checkpoints:TensorFlowCheckpoint", + "tensorflow-checkpoint = git_theta_checkpoints_tensorflow.checkpoints:TensorFlowCheckpoint", + "tf-savedmodel = git_theta_checkpoints_tensorflow.checkpoints:TensorFlowSavedModel", + "tensorflow-savedmodel = git_theta_checkpoints_tensorflow.checkpoints:TensorFlowSavedModel", + ], + "git_theta.plugins.checkpoint.sniffers": [ + "tensorflow = git_theta_checkpoints_tensorflow.sniffer:tensorflow_sniffer", + "tensorflow-savedmodel = git_theta_checkpoints_tensorflow.sniffer:saved_model_sniffer", + ], + }, +) diff --git a/setup.py b/setup.py index 9552c4f..2d7ff1c 100644 --- a/setup.py +++ b/setup.py @@ -43,10 +43,11 @@ def get_version(file_name: str, version_variable: str = "__version__") -> str: # Packages to install for using different deep learning frameworks. frameworks_require = { - "pytorch": ["torch"], - "tensorflow": ["tensorflow"], - "flax": ["flax"], - "safetensors": ["safetensors"], + "pytorch": ["git_theta_checkpoints_pytorch"], + "torch": ["git_theta_checkpoints_pytorch"], + "tensorflow": ["git_theta_checkpoints_tensorflow"], + "flax": ["git_theta_checkpoints_flax"], + "safetensors": ["git_theta_checkpoints_safetensors"], } @@ -113,19 +114,8 @@ def get_version(file_name: str, version_variable: str = "__version__") -> str: "git-theta-merge = git_theta.scripts.git_theta_merge:main", "git-theta-diff = git_theta.scripts.git_theta_diff:main", ], - "git_theta.plugins.checkpoints": [ - "pytorch = git_theta.checkpoints.pickled_dict_checkpoint:PickledDictCheckpoint", - "pickled-dict = git_theta.checkpoints.pickled_dict_checkpoint:PickledDictCheckpoint", - "tf = git_theta.checkpoints.tensorflow_checkpoint:TensorFlowCheckpoint", - "tensorflow = git_theta.checkpoints.tensorflow_checkpoint:TensorFlowCheckpoint", - "tensorflow-checkpoint = git_theta.checkpoints.tensorflow_checkpoint:TensorFlowCheckpoint", - "tf-savedmodel = git_theta.checkpoints.tensorflow_checkpoint:TensorFlowSavedModel", - "tensorflow-savedmodel = git_theta.checkpoints.tensorflow_checkpoint:TensorFlowSavedModel", - "flax = git_theta.checkpoints.flax_checkpoint:FlaxCheckpoint", - "flax-checkpoint = git_theta.checkpoints.flax_checkpoint:FlaxCheckpoint", - "safetensors = git_theta.checkpoints.safetensors_checkpoint:SafeTensorsCheckpoint", - "safetensors-checkpoint = git_theta.checkpoints.safetensors_checkpoint:SafeTensorsCheckpoint", - ], + "git_theta.plugins.checkpoints": [], + "git_theta.plugins.checkpoint.sniffers": [], "git_theta.plugins.updates": [ "dense = git_theta.updates.dense:DenseUpdate", "sparse = git_theta.updates.sparse:SparseUpdate",