add version.py

setup.py

@@ -5,6 +5,11 @@ import torch
 from setuptools import setup, find_packages
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension, ROCM_HOME
 
+from typing import Optional, Union
+import subprocess
+from pathlib import Path
+
+
 # ninja build does not work unless include_dirs are abs path
 this_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -92,9 +97,44 @@ else:
         cuda_ext_helper('fastfold_softmax_cuda', ['softmax_cuda.cpp', 'softmax_cuda_kernel.cu'],
                         extra_cuda_flags + cc_flag))
 
+
+def get_sha(pytorch_root: Union[str, Path]) -> str:
+    try:
+        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=pytorch_root).decode('ascii').strip()
+    except Exception:
+        return 'Unknown'
+
+
+def get_version_add(sha: Optional[str] = None) -> str:
+    add_version_path = "jax/version.py"
+    if sha != 'Unknown':
+        if sha is None:
+            sha_path = os.getenv('FASTFOLD_DOWNLOAD_PATH', "")
+            sha = get_sha(sha_path)
+        version = 'git' + sha[:7]
+
+    if os.getenv('FASTFOLD_BUILD_VERSION'):
+        version_dtk = os.getenv('FASTFOLD_BUILD_VERSION', "")
+        version += "." + version_dtk
+
+    with open(add_version_path, encoding="utf-8", mode="w") as file:
+        file.write("__version__='0.2.21'+'+{}'\n".format(version))
+        file.write("_minimum_jaxlib_version='0.1.69'\n")
+    file.close()
+
+
+def get_version():
+    get_version_add()
+    version_file = 'version.py'
+    with open(version_file, encoding='utf-8') as f:
+        exec(compile(f.read(), version_file, 'exec'))
+    return locals()['__version__']
+
+
 setup(
     name='fastfold',
-    version='0.2.0',
+    #version='0.2.0',
+    version=get_version(),
     packages=find_packages(exclude=(
         'assets',
         'benchmark',

version.py

+__version__ = "0.2.0"
\ No newline at end of file