Add version display on table and JSON output (#36)

Author: alvarobartt

src/hf_mem/__init__.py

@@ -0,0 +1,3 @@
+from importlib.metadata import version
+
+__version__ = version("hf-mem")

src/hf_mem/cli.py

@@ -11,6 +11,7 @@
 
 import httpx
 
+from hf_mem import __version__
 from hf_mem.metadata import parse_safetensors_metadata
 from hf_mem.print import print_report
 from hf_mem.types import TorchDtypes, get_safetensors_dtype_bytes, torch_dtype_to_safetensors_dtype
@@ -364,12 +365,15 @@ async def fetch_with_semaphore(url: str) -> Dict[str, Any]:
                     cache_size *= batch_size
 
     if json_output:
-        out = {"model_id": model_id, "revision": revision, **asdict(metadata)}
+        from hf_mem import __version__
+
+        out = {"version": __version__, "model_id": model_id, "revision": revision, **asdict(metadata)}
         if experimental and cache_size:
             out["max_model_len"] = max_model_len
             out["batch_size"] = batch_size
             out["cache_size"] = cache_size
             out["cache_dtype"] = cache_dtype  # type: ignore
+
         print(json.dumps(out))
     else:
         # TODO: Use a `KvCache` dataclass instead and make sure that the JSON output is aligned

src/hf_mem/print.py

@@ -1,6 +1,7 @@
 import warnings
 from typing import Any, Dict, Literal, Optional
 
+from hf_mem import __version__
 from hf_mem.metadata import SafetensorsMetadata
 
 MIN_NAME_LEN = 5
@@ -179,6 +180,9 @@ def print_report(
         )
     _print_divider(data_col_width + 1, "top")
 
+    _print_row("VERSION", f"hf-mem {__version__}", data_col_width)
+    _print_divider(data_col_width + 1)
+
     if cache:
         total_text = f"{_bytes_to_gib(combined_total):.2f} GiB ({_format_short_number(metadata.param_count)} PARAMS + KV CACHE)"
         total_bar = _make_bar(combined_total, combined_total, data_col_width)