Allow adding an optional local version to the package version

flash_attn/__init__.py

+__version__ = "1.0.5"

setup.py

@@ -2,6 +2,8 @@
 import sys
 import warnings
 import os
+import re
+import ast
 from pathlib import Path
 from packaging.version import parse, Version
 
@@ -160,9 +162,19 @@ ext_modules.append(
     )
 )
 
+def get_package_version():
+    with open(Path(this_dir) / "flash_attn" / "__init__.py", "r") as f:
+        version_match = re.search(r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE)
+    public_version = ast.literal_eval(version_match.group(1))
+    local_version = os.environ.get("FLASH_ATTN_LOCAL_VERSION")
+    if local_version:
+        return f"{public_version}+{local_version}"
+    else:
+        return str(public_version)
+
 setup(
     name="flash_attn",
-    version="1.0.5",
+    version=get_package_version(),
     packages=find_packages(
         exclude=("build", "csrc", "include", "tests", "dist", "docs", "benchmarks", "flash_attn.egg-info",)
     ),