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.

Author: blester125

git_theta/checkpoints/__init__.py

@@ -4,4 +4,5 @@
     Checkpoint,
     get_checkpoint_handler,
     get_checkpoint_handler_name,
+    sniff_checkpoint,
 )

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}")

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)

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)

plugins/checkpoints/flax/git_theta_checkpoints_flax/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

git_theta/checkpoints/flax_checkpoint.py plugins/checkpoints/flax/git_theta_checkpoints_flax/checkpoints.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")

plugins/checkpoints/flax/git_theta_checkpoints_flax/sniffer.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",
+        ],
+    },
+)

plugins/checkpoints/flax/setup.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/__init__.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

git_theta/checkpoints/pickled_dict_checkpoint.py plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/checkpoints.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",
+        ],
+    },
+)

plugins/checkpoints/pytorch/git_theta_checkpoints_pytorch/sniffer.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

plugins/checkpoints/pytorch/setup.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")

plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/__init__.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")

git_theta/checkpoints/safetensors_checkpoint.py plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/checkpoints.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",
+        ],
+    },
+)

plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/flycheck_checkpoints.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

plugins/checkpoints/safetensors/git_theta_checkpoints_safetensors/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