Fix default component and add missing table divider (#10)

* Fix default component name for Transformers

* Run `uv version 0.2.2`

* Fix `print_report` divider on multiple dtypes

* Empty commit to add co-author

Contributed with #9, which is
a similar fix for the `__metadata__` validation for Transformers models
as the one in this PR, hence adding them as co-author for credits,
thanks!

Co-Authored-By: Bae-ChangHyun <[email protected]>

* Run `pre-commit run --all-files`

---------

Co-authored-by: Bae-ChangHyun <[email protected]>

Author: alvarobartt

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hf-mem"
-version = "0.2.1"
+version = "0.2.2"
 description = "A CLI to estimate inference memory requirements for Hugging Face models, written in Python"
 readme = "README.md"
 authors = [

src/hf_mem/cli.py

@@ -99,10 +99,9 @@ async def run(
         # set the default component name to `0_Transformer` as defined in modules.json
         if "config_sentence_transformers.json" in file_paths:
             raw_metadata = {"0_Transformer": raw_metadata}
-
-        # NOTE: If the model is a transformers model, then we simply set the component name to `Transformer`, to
-        # make sure that we provide the expected input to the `parse_safetensors_metadata`
-        if "__metadata__" in raw_metadata:
+        else:
+            # NOTE: If the model is a transformers model, then we simply set the component name to `Transformer`, to
+            # make sure that we provide the expected input to the `parse_safetensors_metadata`
             raw_metadata = {"Transformer": raw_metadata}
 
         metadata = parse_safetensors_metadata(raw_metadata=raw_metadata)
@@ -132,10 +131,9 @@ async def fetch_with_semaphore(url: str) -> Dict[str, Any]:
         # set the default component name to `0_Transformer` as defined in modules.json
         if "config_sentence_transformers.json" in file_paths:
             raw_metadata = {"0_Transformer": raw_metadata}
-
-        # NOTE: If the model is a transformers model, then we simply set the component name to `Transformer`, to
-        # make sure that we provide the expected input to the `parse_safetensors_metadata`
-        if "__metadata__" in raw_metadata:
+        else:
+            # NOTE: If the model is a transformers model, then we simply set the component name to `Transformer`, to
+            # make sure that we provide the expected input to the `parse_safetensors_metadata`
             raw_metadata = {"Transformer": raw_metadata}
 
         metadata = parse_safetensors_metadata(raw_metadata=raw_metadata)

src/hf_mem/print.py

@@ -169,7 +169,7 @@ def print_report(
                 for _, dtype_metadata in value.dtypes.items()
             ]
         )
-        for dtype, dtype_metadata in value.dtypes.items():
+        for idx, (dtype, dtype_metadata) in enumerate(value.dtypes.items()):
             gb_text = (
                 f"{_bytes_to_gb(dtype_metadata.bytes_count):.2f} / {_bytes_to_gb(metadata.bytes_count):.2f} GB"
             )
@@ -190,4 +190,7 @@ def print_report(
                 current_len,
             )
 
+            if idx < len(value.dtypes) - 1:
+                _print_divider(current_len + 1)
+
     _print_divider(current_len + 1, "bottom")