Merge pull request #227 from google-research/rajat_dev

Revising messages during import

.github/workflows/main.yml

@@ -23,7 +23,7 @@ jobs:
       # e.g. poetry version 0.1.${{ github.run_number }}
       - name: Set Version number
         run: |
-          poetry version 1.2.7
+          poetry version 1.2.8
       - name: Build and Publish to PyPI
         run: |
           poetry config pypi-token.pypi ${{ secrets.PYPI_API_TOKEN }}

pyproject.toml

@@ -2,7 +2,7 @@
 name = "timesfm"
 packages = [{ include = "timesfm", from = "src" }]
 description = "Open weights time-series foundation model from Google Research."
-version = "1.2.7"
+version = "1.2.8"
 authors = [
     "Rajat Sen <senrajat@google.com>",
     "Yichen Zhou <yichenzhou@google.com>",

src/timesfm/__init__.py

@@ -12,14 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """TimesFM init file."""
+
 print(
-    "TimesFM v1.2.0. See https://github.com/google-research/timesfm/blob/master/README.md for updated APIs."
+    " See https://github.com/google-research/timesfm/blob/master/README.md for updated APIs."
 )
-from timesfm.timesfm_base import freq_map, TimesFmCheckpoint, TimesFmHparams, TimesFmBase
+from timesfm.timesfm_base import (
+    freq_map,
+    TimesFmCheckpoint,
+    TimesFmHparams,
+    TimesFmBase,
+)
+import sys
+
 try:
-  print("Loaded Jax TimesFM.")
     from timesfm.timesfm_jax import TimesFmJax as TimesFm
     from timesfm import data_loader
+
+    print(f"Loaded Jax TimesFM, likely because python version is {sys.version}.")
 except Exception as _:
-  print("Loaded PyTorch TimesFM.")
     from timesfm.timesfm_torch import TimesFmTorch as TimesFm
+
+    print(f"Loaded PyTorch TimesFM, likely because python version is {sys.version}.")