From c9d09cea11a68dfa1c822f3bb260440165c2a599 Mon Sep 17 00:00:00 2001
From: Eitan Turok <150733043+eitanturok@users.noreply.github.com>
Date: Wed, 31 Jul 2024 16:40:47 -0400
Subject: [PATCH] Read Package Version Better (#1415)

Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
---
 llmfoundry/__init__.py                      |  4 ++--
 llmfoundry/_version.py                      |  6 ++++++
 llmfoundry/callbacks/async_eval_callback.py |  3 ++-
 setup.py                                    | 18 ++++++++----------
 4 files changed, 18 insertions(+), 13 deletions(-)
 create mode 100644 llmfoundry/_version.py

diff --git a/llmfoundry/__init__.py b/llmfoundry/__init__.py
index 8dbd180c0a..b851aaa559 100644
--- a/llmfoundry/__init__.py
+++ b/llmfoundry/__init__.py
@@ -50,6 +50,7 @@
     tokenizers,
     utils,
 )
+from llmfoundry._version import __version__
 from llmfoundry.data import StreamingFinetuningDataset, StreamingTextDataset
 from llmfoundry.eval import InContextLearningDataset, InContextLearningMetric
 from llmfoundry.models.hf import ComposerHFCausalLM
@@ -63,6 +64,7 @@
 from llmfoundry.optim import DecoupledLionW
 
 __all__ = [
+    '__version__',
     'StreamingFinetuningDataset',
     'StreamingTextDataset',
     'InContextLearningDataset',
@@ -87,5 +89,3 @@
     'tokenizers',
     'utils',
 ]
-
-__version__ = '0.11.0.dev0'
diff --git a/llmfoundry/_version.py b/llmfoundry/_version.py
new file mode 100644
index 0000000000..4c11746b43
--- /dev/null
+++ b/llmfoundry/_version.py
@@ -0,0 +1,6 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""The LLM Foundry Version."""
+
+__version__ = '0.11.0.dev'
diff --git a/llmfoundry/callbacks/async_eval_callback.py b/llmfoundry/callbacks/async_eval_callback.py
index 646d86c8d3..1b3c31e861 100644
--- a/llmfoundry/callbacks/async_eval_callback.py
+++ b/llmfoundry/callbacks/async_eval_callback.py
@@ -557,7 +557,8 @@ def launch_run(self, checkpoint: str, current_interval: Time) -> Run:
                 installation_path = i['path']
 
         if not found_llm_foundry:
-            from llmfoundry import __version__ as latest_foundry_version
+            from llmfoundry._version import \
+                __version__ as latest_foundry_version
 
             # If github integration is not found, foundry is likely installed
             # through the run command. In this case, we'll add the integration
diff --git a/setup.py b/setup.py
index bc84f9fac8..185d3970f7 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 import copy
 import os
-import re
+from typing import Any, Dict, Mapping
 
 import setuptools
 from setuptools import setup
@@ -15,17 +15,15 @@
 _REPO_REAL_PATH = os.path.dirname(os.path.realpath(__file__))
 _PACKAGE_REAL_PATH = os.path.join(_REPO_REAL_PATH, _PACKAGE_DIR)
 
-# Read the repo version
+# Read the llm-foundry version
 # We can't use `.__version__` from the library since it's not installed yet
-with open(os.path.join(_PACKAGE_REAL_PATH, '__init__.py')) as f:
+version_path = os.path.join(_PACKAGE_REAL_PATH, '_version.py')
+with open(version_path, encoding='utf-8') as f:
+    version_globals: Dict[str, Any] = {}
+    version_locals: Mapping[str, object] = {}
     content = f.read()
-# regex: '__version__', whitespace?, '=', whitespace, quote, version, quote
-# we put parens around the version so that it becomes elem 1 of the match
-expr = re.compile(
-    r"""^__version__\s*=\s*['"]([0-9]+\.[0-9]+\.[0-9]+(?:\.\w+)?)['"]""",
-    re.MULTILINE,
-)
-repo_version = expr.findall(content)[0]
+    exec(content, version_globals, version_locals)
+    repo_version = str(version_locals['__version__'])
 
 # Use repo README for PyPi description
 with open('README.md', 'r', encoding='utf-8') as fh: