Finish refactor

.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

src/app.py

@@ -1,8 +1,7 @@
 import gradio as gr
 import pandas as pd
-
-from .hub_utils import check_for_discussion, report_results
-from .model_utils import calculate_memory, get_model
+from hub_utils import check_for_discussion, report_results
+from model_utils import calculate_memory, 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

src/hub_utils.py

@@ -1,33 +1,16 @@
 # Utilities related to searching and posting on the Hub
 import os
 import webbrowser
-from urllib.parse import urlparse
 
 import pandas as pd
 from huggingface_hub import HfApi
-
-from .model_utils import calculate_memory, get_model
-
-
-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:]
+from model_utils import calculate_memory, extract_from_url, get_model
 
 
 def check_for_discussion(model_name: str):
     "Checks if an automated discussion has been opened on the model by `model-sizer-bot`"
     api = HfApi(token=os.environ.get("HUGGINGFACE_API_LOGIN", None))
+    model_name = extract_from_url(model_name)
     discussions = list(api.get_repo_discussions(model_name))
     return any(
         discussion.title == "[AUTOMATED] Model Memory Requirements" and discussion.author == "model-sizer-bot"
@@ -38,13 +21,12 @@ def check_for_discussion(model_name: str):
 def report_results(model_name, library, access_token):
     "Reports the results of a memory calculation to the model's discussion page, and opens a new tab to it afterwards"
     model = get_model(model_name, library, access_token)
-    data = calculate_memory(model, ["fp32", "fp16", "int8", "int4"])
-    minimum = data[0]
-    data = pd.DataFrame(data).to_markdown(index=False)
+    data = calculate_memory(model, ["float32", "float16/bfloat16", "int8", "int4"])
+    df = pd.DataFrame(data).to_markdown(index=False)
 
     post = f"""# Model Memory Requirements\n
 
-You will need about {minimum[1]} VRAM to load this model for inference, and {minimum[3]} VRAM to train it using Adam.
+You will need about {data[1]} VRAM to load this model for inference, and {data[3]} VRAM to train it using Adam.
     
 These calculations were measured from the [Model Memory Utility Space](https://hf.co/spaces/hf-accelerate/model-memory-utility) on the Hub.
     
@@ -55,7 +37,7 @@ When training with `Adam`, you can expect roughly 4x the reported results to be
 
 ## Results:
 
-{data}
+{df}
 """
     api = HfApi(token=os.environ.get("HUGGINGFACE_API_LOGIN", None))
     discussion = api.create_discussion(model_name, "[AUTOMATED] Model Memory Requirements", description=post)

src/model_utils.py

@@ -1,16 +1,32 @@
 # 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
 
-from .hub_utils import extract_from_url
-
 
 DTYPE_MODIFIER = {"float32": 1, "float16/bfloat16": 2, "int8": 4, "int4": 8}
 
 
+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"):