Add initial version

.gitignore

@@ -0,0 +1,144 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# VSCode
+.vscode
+
+# IntelliJ
+.idea
+
+# Mac .DS_Store
+.DS_Store
+
+# More test things
+wandb
+
+# ruff
+.ruff_cache

README.md

@@ -1,5 +1,5 @@
 ---
-title: Model Visualizer
+title: Model Representation
 emoji: 👁
 colorFrom: gray
 colorTo: red

src/app.py

@@ -0,0 +1,42 @@
+import gradio as gr
+import pandas as pd
+from huggingface_hub.utils import HfHubHTTPError
+from accelerate.commands.estimate import check_has_model, create_empty_model
+from accelerate.utils import calculate_maximum_sizes
+from model_utils import get_model
+
+# We need to store them as globals because gradio doesn't have a way for us to pass them in to the button
+MODEL = None
+
+
+def get_results(model_name: str, library: str, precision: list, training: list,  access_token: str, zero_stage: int, num_nodes: int, num_gpus: int, offloading: list, zero_init: list, additional_buffer_factor: float):
+    global MODEL
+    
+    MODEL = get_model(model_name, library, access_token)
+
+    data = MODEL.__repr__()
+    
+    title = f"## Model Representation for '{model_name}'"
+    return [title, gr.update(visible=True, value=data)]
+
+with gr.Blocks() as demo:
+    with gr.Column():
+        out_text = gr.Markdown()
+        out = gr.Code()
+
+        with gr.Row():
+            inp = gr.Textbox(label="Model Name or URL", value="bert-base-cased")
+        with gr.Row():
+            library = gr.Radio(["auto", "transformers", "timm"], label="Library", value="auto")
+            access_token = gr.Textbox(label="API Token", placeholder="Optional (for gated models)")
+        with gr.Row():
+            btn = gr.Button("Calculate Memory Usage")
+        
+
+    btn.click(
+        get_results,
+        inputs=[inp, library, access_token,],
+        outputs=[out_text, out],
+    )
+
+demo.launch()

src/model_utils.py

@@ -0,0 +1,76 @@
+# Utilities related to loading in and working with models/specific models
+from urllib.parse import urlparse
+
+import gradio as gr
+import torch
+from accelerate.commands.estimate import check_has_model, create_empty_model
+from accelerate.utils import calculate_maximum_sizes, convert_bytes
+from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
+
+
+DTYPE_MODIFIER = {"float32": 1, "float16/bfloat16": 2, "int8": 4, "int4": 8}
+PRECISION = {"Mixed precision": "mixed", "Single precision": "single"}
+DTYPE = {"float32": torch.float32, "float16/bfloat16": torch.float16}
+
+
+def extract_from_url(name: str):
+    "Checks if `name` is a URL, and if so converts it to a model name"
+    is_url = False
+    try:
+        result = urlparse(name)
+        is_url = all([result.scheme, result.netloc])
+    except Exception:
+        is_url = False
+    # Pass through if not a URL
+    if not is_url:
+        return name
+    else:
+        path = result.path
+        return path[1:]
+
+
+def translate_llama2(text):
+    "Translates llama-2 to its hf counterpart"
+    if not text.endswith("-hf"):
+        return text + "-hf"
+    return text
+
+
+def get_model(model_name: str, library: str, access_token: str):
+    "Finds and grabs model from the Hub, and initializes on `meta`"
+    if "meta-llama" in model_name:
+        model_name = translate_llama2(model_name)
+    if library == "auto":
+        library = None
+    model_name = extract_from_url(model_name)
+    try:
+        model = create_empty_model(model_name, library_name=library, trust_remote_code=True, access_token=access_token)
+    except GatedRepoError:
+        raise gr.Error(
+            f"Model `{model_name}` is a gated model, please ensure to pass in your access token and try again if you have access. You can find your access token here : https://huggingface.co/settings/tokens. "
+        )
+    except RepositoryNotFoundError:
+        raise gr.Error(f"Model `{model_name}` was not found on the Hub, please try another model name.")
+    except ValueError:
+        raise gr.Error(
+            f"Model `{model_name}` does not have any library metadata on the Hub, please manually select a library_name to use (such as `transformers`)"
+        )
+    except (RuntimeError, OSError) as e:
+        library = check_has_model(e)
+        if library != "unknown":
+            raise gr.Error(
+                f"Tried to load `{model_name}` with `{library}` but a possible model to load was not found inside the repo."
+            )
+        raise gr.Error(
+            f"Model `{model_name}` had an error, please open a discussion on the model's page with the error message and name: `{e}`"
+        )
+    except ImportError:
+        # hacky way to check if it works with `trust_remote_code=False`
+        model = create_empty_model(
+            model_name, library_name=library, trust_remote_code=False, access_token=access_token
+        )
+    except Exception as e:
+        raise gr.Error(
+            f"Model `{model_name}` had an error, please open a discussion on the model's page with the error message and name: `{e}`"
+        )
+    return model