mirror of
https://github.com/Laurent2916/model-explorer-refiners.git
synced 2024-09-16 17:45:29 +00:00
initial commit
This commit is contained in:
Laurent
commit
8eacdfda7b
163
.gitignore
vendored
Normal file
163
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
# https://github.com/github/gitignore/blob/main/Python.gitignore
|
||||
# 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/
|
||||
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/
|
||||
cover/
|
||||
|
||||
# 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
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .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
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
|
||||
.pdm.toml
|
||||
.pdm-python
|
||||
.pdm-build/
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__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/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
1
.python-version
Normal file
1
.python-version
Normal file
|
|
@ -0,0 +1 @@
|
|||
3.12.1
|
||||
20
README.md
Normal file
20
README.md
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
# Refiners Model Explorer
|
||||
|
||||
A Model Explorer extension to visualize and explore Refiners-based model.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install git+https://github.com/finegrain-ai/model_explorer_refiners.git
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from refiners.foundationals.dinov2 import DINOv2_small_reg
|
||||
from model_explorer_refiners import RefinersAdapter
|
||||
|
||||
model = DINOv2_small_reg()
|
||||
RefinersAdapter.visualize(model)
|
||||
|
||||
```
|
||||
60
pyproject.toml
Normal file
60
pyproject.toml
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
[project]
|
||||
version = "0.1.0"
|
||||
readme = "README.md"
|
||||
requires-python = ">= 3.9"
|
||||
name = "model-explorer-refiners"
|
||||
description = "Refiners adapter for Model Explorer"
|
||||
authors = [{ name = "Laurent", email = "laurent@lagon.tech" }]
|
||||
dependencies = [
|
||||
"ai-edge-model-explorer",
|
||||
"refiners @ git+https://github.com/finegrain-ai/refiners",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.rye]
|
||||
managed = true
|
||||
dev-dependencies = []
|
||||
|
||||
[tool.hatch.metadata]
|
||||
allow-direct-references = true
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/model_explorer_refiners"]
|
||||
|
||||
[tool.ruff]
|
||||
src = ["src"]
|
||||
line-length = 120
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"A", # flake8-builtins
|
||||
"B", # flake8-bugbear
|
||||
"C90", # mccabe
|
||||
"COM", # flake8-commas
|
||||
"EM", # flake8-errmsg
|
||||
"E", # pycodestyle errors
|
||||
"F", # Pyflakes
|
||||
"G", # flake8-logging-format
|
||||
"I", # isort
|
||||
"N", # pep8-naming
|
||||
"PIE", # flake8-pie
|
||||
"PTH", # flake8-use-pathlib
|
||||
"TD", # flake8-todo
|
||||
"FIX", # flake8-fixme
|
||||
"RUF", # ruff
|
||||
"S", # flake8-bandit
|
||||
"TCH", # flake8-type-checking
|
||||
"TID", # flake8-tidy-imports
|
||||
"UP", # pyupgrade
|
||||
"W", # pycodestyle warnings
|
||||
]
|
||||
|
||||
|
||||
[tool.pyright]
|
||||
include = ["model_explorer_refiners"]
|
||||
reportMissingTypeStubs = false
|
||||
pythonVersion = "3.12"
|
||||
pythonPlatform = "Linux"
|
||||
140
requirements-dev.lock
Normal file
140
requirements-dev.lock
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
# generated by rye
|
||||
# use `rye lock` or `rye sync` to update this lockfile
|
||||
#
|
||||
# last locked with the following flags:
|
||||
# pre: false
|
||||
# features: []
|
||||
# all-features: true
|
||||
# with-sources: false
|
||||
# generate-hashes: false
|
||||
|
||||
-e file:.
|
||||
ai-edge-model-explorer==0.1.7
|
||||
# via model-explorer-refiners
|
||||
ai-edge-model-explorer-adapter==0.1.3
|
||||
# via ai-edge-model-explorer
|
||||
asttokens==2.4.1
|
||||
# via stack-data
|
||||
blinker==1.8.2
|
||||
# via flask
|
||||
certifi==2024.6.2
|
||||
# via requests
|
||||
charset-normalizer==3.3.2
|
||||
# via requests
|
||||
click==8.1.7
|
||||
# via flask
|
||||
decorator==5.1.1
|
||||
# via ipython
|
||||
executing==2.0.1
|
||||
# via stack-data
|
||||
filelock==3.15.4
|
||||
# via torch
|
||||
flask==3.0.3
|
||||
# via ai-edge-model-explorer
|
||||
fsspec==2024.6.1
|
||||
# via torch
|
||||
idna==3.7
|
||||
# via requests
|
||||
ipython==8.26.0
|
||||
# via ai-edge-model-explorer
|
||||
itsdangerous==2.2.0
|
||||
# via flask
|
||||
jaxtyping==0.2.31
|
||||
# via refiners
|
||||
jedi==0.19.1
|
||||
# via ipython
|
||||
jinja2==3.1.4
|
||||
# via flask
|
||||
# via torch
|
||||
markupsafe==2.1.5
|
||||
# via jinja2
|
||||
# via werkzeug
|
||||
matplotlib-inline==0.1.7
|
||||
# via ipython
|
||||
mpmath==1.3.0
|
||||
# via sympy
|
||||
networkx==3.3
|
||||
# via torch
|
||||
numpy==2.0.0
|
||||
# via ai-edge-model-explorer
|
||||
# via refiners
|
||||
nvidia-cublas-cu12==12.1.3.1
|
||||
# via nvidia-cudnn-cu12
|
||||
# via nvidia-cusolver-cu12
|
||||
# via torch
|
||||
nvidia-cuda-cupti-cu12==12.1.105
|
||||
# via torch
|
||||
nvidia-cuda-nvrtc-cu12==12.1.105
|
||||
# via torch
|
||||
nvidia-cuda-runtime-cu12==12.1.105
|
||||
# via torch
|
||||
nvidia-cudnn-cu12==8.9.2.26
|
||||
# via torch
|
||||
nvidia-cufft-cu12==11.0.2.54
|
||||
# via torch
|
||||
nvidia-curand-cu12==10.3.2.106
|
||||
# via torch
|
||||
nvidia-cusolver-cu12==11.4.5.107
|
||||
# via torch
|
||||
nvidia-cusparse-cu12==12.1.0.106
|
||||
# via nvidia-cusolver-cu12
|
||||
# via torch
|
||||
nvidia-nccl-cu12==2.20.5
|
||||
# via torch
|
||||
nvidia-nvjitlink-cu12==12.5.40
|
||||
# via nvidia-cusolver-cu12
|
||||
# via nvidia-cusparse-cu12
|
||||
nvidia-nvtx-cu12==12.1.105
|
||||
# via torch
|
||||
packaging==24.1
|
||||
# via ai-edge-model-explorer
|
||||
# via refiners
|
||||
parso==0.8.4
|
||||
# via jedi
|
||||
pexpect==4.9.0
|
||||
# via ipython
|
||||
pillow==10.3.0
|
||||
# via refiners
|
||||
portpicker==1.6.0
|
||||
# via ai-edge-model-explorer
|
||||
prompt-toolkit==3.0.47
|
||||
# via ipython
|
||||
psutil==6.0.0
|
||||
# via portpicker
|
||||
ptyprocess==0.7.0
|
||||
# via pexpect
|
||||
pure-eval==0.2.2
|
||||
# via stack-data
|
||||
pygments==2.18.0
|
||||
# via ipython
|
||||
refiners @ git+https://github.com/finegrain-ai/refiners@e091788b885c5024b82e18cef082a922cd050481
|
||||
# via model-explorer-refiners
|
||||
requests==2.32.3
|
||||
# via ai-edge-model-explorer
|
||||
safetensors==0.4.3
|
||||
# via refiners
|
||||
six==1.16.0
|
||||
# via asttokens
|
||||
stack-data==0.6.3
|
||||
# via ipython
|
||||
sympy==1.12.1
|
||||
# via torch
|
||||
termcolor==2.4.0
|
||||
# via ai-edge-model-explorer
|
||||
torch==2.3.1
|
||||
# via ai-edge-model-explorer
|
||||
# via refiners
|
||||
traitlets==5.14.3
|
||||
# via ipython
|
||||
# via matplotlib-inline
|
||||
typeguard==2.13.3
|
||||
# via jaxtyping
|
||||
typing-extensions==4.12.2
|
||||
# via ai-edge-model-explorer
|
||||
# via torch
|
||||
urllib3==2.2.2
|
||||
# via requests
|
||||
wcwidth==0.2.13
|
||||
# via prompt-toolkit
|
||||
werkzeug==3.0.3
|
||||
# via flask
|
||||
140
requirements.lock
Normal file
140
requirements.lock
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
# generated by rye
|
||||
# use `rye lock` or `rye sync` to update this lockfile
|
||||
#
|
||||
# last locked with the following flags:
|
||||
# pre: false
|
||||
# features: []
|
||||
# all-features: true
|
||||
# with-sources: false
|
||||
# generate-hashes: false
|
||||
|
||||
-e file:.
|
||||
ai-edge-model-explorer==0.1.7
|
||||
# via model-explorer-refiners
|
||||
ai-edge-model-explorer-adapter==0.1.3
|
||||
# via ai-edge-model-explorer
|
||||
asttokens==2.4.1
|
||||
# via stack-data
|
||||
blinker==1.8.2
|
||||
# via flask
|
||||
certifi==2024.6.2
|
||||
# via requests
|
||||
charset-normalizer==3.3.2
|
||||
# via requests
|
||||
click==8.1.7
|
||||
# via flask
|
||||
decorator==5.1.1
|
||||
# via ipython
|
||||
executing==2.0.1
|
||||
# via stack-data
|
||||
filelock==3.15.4
|
||||
# via torch
|
||||
flask==3.0.3
|
||||
# via ai-edge-model-explorer
|
||||
fsspec==2024.6.1
|
||||
# via torch
|
||||
idna==3.7
|
||||
# via requests
|
||||
ipython==8.26.0
|
||||
# via ai-edge-model-explorer
|
||||
itsdangerous==2.2.0
|
||||
# via flask
|
||||
jaxtyping==0.2.31
|
||||
# via refiners
|
||||
jedi==0.19.1
|
||||
# via ipython
|
||||
jinja2==3.1.4
|
||||
# via flask
|
||||
# via torch
|
||||
markupsafe==2.1.5
|
||||
# via jinja2
|
||||
# via werkzeug
|
||||
matplotlib-inline==0.1.7
|
||||
# via ipython
|
||||
mpmath==1.3.0
|
||||
# via sympy
|
||||
networkx==3.3
|
||||
# via torch
|
||||
numpy==2.0.0
|
||||
# via ai-edge-model-explorer
|
||||
# via refiners
|
||||
nvidia-cublas-cu12==12.1.3.1
|
||||
# via nvidia-cudnn-cu12
|
||||
# via nvidia-cusolver-cu12
|
||||
# via torch
|
||||
nvidia-cuda-cupti-cu12==12.1.105
|
||||
# via torch
|
||||
nvidia-cuda-nvrtc-cu12==12.1.105
|
||||
# via torch
|
||||
nvidia-cuda-runtime-cu12==12.1.105
|
||||
# via torch
|
||||
nvidia-cudnn-cu12==8.9.2.26
|
||||
# via torch
|
||||
nvidia-cufft-cu12==11.0.2.54
|
||||
# via torch
|
||||
nvidia-curand-cu12==10.3.2.106
|
||||
# via torch
|
||||
nvidia-cusolver-cu12==11.4.5.107
|
||||
# via torch
|
||||
nvidia-cusparse-cu12==12.1.0.106
|
||||
# via nvidia-cusolver-cu12
|
||||
# via torch
|
||||
nvidia-nccl-cu12==2.20.5
|
||||
# via torch
|
||||
nvidia-nvjitlink-cu12==12.5.40
|
||||
# via nvidia-cusolver-cu12
|
||||
# via nvidia-cusparse-cu12
|
||||
nvidia-nvtx-cu12==12.1.105
|
||||
# via torch
|
||||
packaging==24.1
|
||||
# via ai-edge-model-explorer
|
||||
# via refiners
|
||||
parso==0.8.4
|
||||
# via jedi
|
||||
pexpect==4.9.0
|
||||
# via ipython
|
||||
pillow==10.3.0
|
||||
# via refiners
|
||||
portpicker==1.6.0
|
||||
# via ai-edge-model-explorer
|
||||
prompt-toolkit==3.0.47
|
||||
# via ipython
|
||||
psutil==6.0.0
|
||||
# via portpicker
|
||||
ptyprocess==0.7.0
|
||||
# via pexpect
|
||||
pure-eval==0.2.2
|
||||
# via stack-data
|
||||
pygments==2.18.0
|
||||
# via ipython
|
||||
refiners @ git+https://github.com/finegrain-ai/refiners@e091788b885c5024b82e18cef082a922cd050481
|
||||
# via model-explorer-refiners
|
||||
requests==2.32.3
|
||||
# via ai-edge-model-explorer
|
||||
safetensors==0.4.3
|
||||
# via refiners
|
||||
six==1.16.0
|
||||
# via asttokens
|
||||
stack-data==0.6.3
|
||||
# via ipython
|
||||
sympy==1.12.1
|
||||
# via torch
|
||||
termcolor==2.4.0
|
||||
# via ai-edge-model-explorer
|
||||
torch==2.3.1
|
||||
# via ai-edge-model-explorer
|
||||
# via refiners
|
||||
traitlets==5.14.3
|
||||
# via ipython
|
||||
# via matplotlib-inline
|
||||
typeguard==2.13.3
|
||||
# via jaxtyping
|
||||
typing-extensions==4.12.2
|
||||
# via ai-edge-model-explorer
|
||||
# via torch
|
||||
urllib3==2.2.2
|
||||
# via requests
|
||||
wcwidth==0.2.13
|
||||
# via prompt-toolkit
|
||||
werkzeug==3.0.3
|
||||
# via flask
|
||||
3
src/model_explorer_refiners/__init__.py
Normal file
3
src/model_explorer_refiners/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from model_explorer_refiners.main import RefinersAdapter
|
||||
|
||||
__all__ = ["RefinersAdapter"]
|
||||
298
src/model_explorer_refiners/convert.py
Normal file
298
src/model_explorer_refiners/convert.py
Normal file
|
|
@ -0,0 +1,298 @@
|
|||
import refiners.fluxion.layers as fl
|
||||
from model_explorer import ModelExplorerGraphs
|
||||
from model_explorer.graph_builder import Graph, GraphNode, IncomingEdge, KeyValue
|
||||
|
||||
from model_explorer_refiners.utils import find_node
|
||||
|
||||
|
||||
def convert_chain(
|
||||
chain: fl.Chain,
|
||||
input_nodes: list[GraphNode],
|
||||
) -> tuple[list[GraphNode], list[GraphNode]]:
|
||||
"""Convert a refiners chain layer to nodes.
|
||||
|
||||
Args:
|
||||
chain: the chain layer to convert.
|
||||
input_nodes: the input nodes of the chain.
|
||||
"""
|
||||
nodes: list[GraphNode] = []
|
||||
previous_nodes = input_nodes
|
||||
for layer, parent_layer in chain.walk(recurse=False):
|
||||
module_nodes, previous_nodes = convert_module(layer, parent_layer, previous_nodes)
|
||||
nodes.extend(module_nodes)
|
||||
|
||||
return nodes, previous_nodes
|
||||
|
||||
|
||||
def convert_passthrough(
|
||||
passthrough: fl.Passthrough,
|
||||
input_nodes: list[GraphNode],
|
||||
) -> tuple[list[GraphNode], list[GraphNode]]:
|
||||
"""Convert a refiners passthrough layer to nodes.
|
||||
|
||||
Args:
|
||||
passthrough: the passthrough layer to convert.
|
||||
input_nodes: the input nodes of the passthrough.
|
||||
"""
|
||||
nodes, _ = convert_chain(passthrough, input_nodes)
|
||||
return nodes, input_nodes
|
||||
|
||||
|
||||
def convert_residual(
|
||||
residual: fl.Residual,
|
||||
input_nodes: list[GraphNode],
|
||||
) -> tuple[list[GraphNode], list[GraphNode]]:
|
||||
"""Convert a refiners residual layer to nodes.
|
||||
|
||||
Args:
|
||||
residual: the residual layer to convert.
|
||||
input_nodes: the input nodes of the residual.
|
||||
"""
|
||||
nodes: list[GraphNode] = []
|
||||
previous_nodes = input_nodes
|
||||
for layer, parent_layer in residual.walk(recurse=False):
|
||||
module_nodes, previous_nodes = convert_module(layer, parent_layer, previous_nodes)
|
||||
nodes.extend(module_nodes)
|
||||
|
||||
# add a summation node
|
||||
summation_node = GraphNode(
|
||||
namespace=residual.get_path().replace(".", "/"),
|
||||
id=residual.get_path(),
|
||||
label="+",
|
||||
)
|
||||
for node in previous_nodes:
|
||||
summation_node.incomingEdges.append(
|
||||
IncomingEdge(
|
||||
sourceNodeId=node.id,
|
||||
),
|
||||
)
|
||||
for input_node in input_nodes:
|
||||
summation_node.incomingEdges.append(
|
||||
IncomingEdge(
|
||||
sourceNodeId=input_node.id,
|
||||
),
|
||||
)
|
||||
nodes.append(summation_node)
|
||||
|
||||
return nodes, [summation_node]
|
||||
|
||||
|
||||
def convert_parallel(
|
||||
parallel: fl.Parallel,
|
||||
input_nodes: list[GraphNode],
|
||||
) -> tuple[list[GraphNode], list[GraphNode]]:
|
||||
"""Convert a refiners parallel layer to nodes.
|
||||
|
||||
Args:
|
||||
parallel: the parallel layer to convert.
|
||||
input_nodes: the input nodes of the parallel layer.
|
||||
"""
|
||||
nodes: list[GraphNode] = []
|
||||
output_nodes: list[GraphNode] = []
|
||||
for layer, parent_layer in parallel.walk(recurse=False):
|
||||
module_nodes, module_output_nodes = convert_module(layer, parent_layer, input_nodes)
|
||||
nodes.extend(module_nodes)
|
||||
output_nodes.extend(module_output_nodes)
|
||||
|
||||
return nodes, output_nodes
|
||||
|
||||
|
||||
def convert_concatenate(
|
||||
concatenate: fl.Concatenate,
|
||||
input_nodes: list[GraphNode],
|
||||
) -> tuple[list[GraphNode], list[GraphNode]]:
|
||||
"""Convert a refiners concatenate layer to nodes.
|
||||
|
||||
Args:
|
||||
concatenate: the concatenate layer to convert.
|
||||
input_nodes: the input nodes of the concatenate layer.
|
||||
"""
|
||||
nodes: list[GraphNode] = []
|
||||
output_nodes: list[GraphNode] = []
|
||||
for layer, parent_layer in concatenate.walk(recurse=False):
|
||||
module_nodes, module_output_nodes = convert_module(layer, parent_layer, input_nodes)
|
||||
nodes.extend(module_nodes)
|
||||
output_nodes.extend(module_output_nodes)
|
||||
|
||||
# add a concatenation node
|
||||
concat_node = GraphNode(
|
||||
namespace=concatenate.get_path().replace(".", "/"),
|
||||
id=concatenate.get_path(),
|
||||
label="+",
|
||||
)
|
||||
for node in output_nodes:
|
||||
concat_node.incomingEdges.append(
|
||||
IncomingEdge(
|
||||
sourceNodeId=node.id,
|
||||
),
|
||||
)
|
||||
nodes.append(concat_node)
|
||||
|
||||
return nodes, [concat_node]
|
||||
|
||||
|
||||
def convert_distribute(
|
||||
distribute: fl.Distribute,
|
||||
input_nodes: list[GraphNode],
|
||||
) -> tuple[list[GraphNode], list[GraphNode]]:
|
||||
"""Convert a refiners distribute layer to nodes.
|
||||
|
||||
Args:
|
||||
distribute: the distribute layer to convert.
|
||||
input_nodes: the input nodes of the distribute layer.
|
||||
"""
|
||||
nodes: list[GraphNode] = []
|
||||
output_nodes: list[GraphNode] = []
|
||||
for (layer, parent_layer), input_node in zip(distribute.walk(recurse=False), input_nodes, strict=True):
|
||||
module_nodes, previous_nodes = convert_module(layer, parent_layer, [input_node])
|
||||
nodes.extend(module_nodes)
|
||||
output_nodes.extend(previous_nodes)
|
||||
|
||||
return nodes, output_nodes
|
||||
|
||||
|
||||
def convert_other(
|
||||
module: fl.Module,
|
||||
parent_module: fl.Chain | None,
|
||||
input_nodes: list[GraphNode],
|
||||
) -> tuple[list[GraphNode], list[GraphNode]]:
|
||||
"""Convert a refiners module to nodes.
|
||||
|
||||
Args:
|
||||
module: the module to convert.
|
||||
parent_module: the parent module of the module.
|
||||
input_nodes: the input nodes of the module.
|
||||
"""
|
||||
node = GraphNode(
|
||||
namespace=parent_module.get_path().replace(".", "/") if parent_module else "",
|
||||
id=module.get_path(parent_module),
|
||||
label=module._get_name(), # type: ignore
|
||||
attrs=[
|
||||
KeyValue(
|
||||
key=key,
|
||||
value=str(value),
|
||||
)
|
||||
for key, value in module.basic_attributes().items()
|
||||
],
|
||||
)
|
||||
setattr(node, "refiners_module", module) # noqa: B010
|
||||
for input_node in input_nodes:
|
||||
node.incomingEdges.append(
|
||||
IncomingEdge(
|
||||
sourceNodeId=input_node.id,
|
||||
),
|
||||
)
|
||||
return [node], [node]
|
||||
|
||||
|
||||
def convert_module( # type: ignore
|
||||
module: fl.Module,
|
||||
parent_module: fl.Chain | None,
|
||||
input_nodes: list[GraphNode],
|
||||
) -> tuple[list[GraphNode], list[GraphNode]]:
|
||||
"""Convert a refiners module to nodes.
|
||||
|
||||
Args:
|
||||
module: the module to convert.
|
||||
parent_module: the parent module of the module.
|
||||
input_nodes: the input nodes of the module.
|
||||
"""
|
||||
match module:
|
||||
case fl.Parallel():
|
||||
return convert_parallel(module, input_nodes)
|
||||
case fl.Distribute():
|
||||
return convert_distribute(module, input_nodes)
|
||||
case fl.Concatenate():
|
||||
return convert_concatenate(module, input_nodes)
|
||||
case fl.Residual():
|
||||
return convert_residual(module, input_nodes)
|
||||
case fl.Passthrough():
|
||||
return convert_passthrough(module, input_nodes)
|
||||
case fl.Chain():
|
||||
return convert_chain(module, input_nodes)
|
||||
case fl.UseContext():
|
||||
return convert_other(module, parent_module=parent_module, input_nodes=[])
|
||||
case fl.Parameter():
|
||||
return convert_other(module, parent_module=parent_module, input_nodes=[])
|
||||
case _:
|
||||
return convert_other(module, parent_module, input_nodes)
|
||||
|
||||
|
||||
def convert_model(model: fl.Chain) -> ModelExplorerGraphs:
|
||||
"""Convert a refiners model to a Graph.
|
||||
|
||||
Args:
|
||||
model: the model to convert.
|
||||
"""
|
||||
# initialize the graph, with an input and output node
|
||||
graph = Graph(id=model._get_name()) # type: ignore
|
||||
input_node = GraphNode(id="input", label="input")
|
||||
graph.nodes.append(input_node)
|
||||
output_node = GraphNode(id="output", label="output")
|
||||
graph.nodes.append(output_node)
|
||||
|
||||
# convert the model
|
||||
model_nodes, model_output_nodes = convert_module(
|
||||
module=model,
|
||||
parent_module=None,
|
||||
input_nodes=[input_node],
|
||||
)
|
||||
graph.nodes.extend(model_nodes)
|
||||
|
||||
# connect the model's output to the output node
|
||||
for node in model_output_nodes:
|
||||
output_node.incomingEdges.append(
|
||||
IncomingEdge(
|
||||
sourceNodeId=node.id,
|
||||
),
|
||||
)
|
||||
|
||||
# connect the model's context nodes
|
||||
for node in graph.nodes:
|
||||
refiners_module = getattr(node, "refiners_module", None)
|
||||
if isinstance(refiners_module, fl.SetContext):
|
||||
# find the context node
|
||||
context_node = find_node(
|
||||
graph=graph,
|
||||
node_id=f"context.{refiners_module.context}.{refiners_module.key}",
|
||||
)
|
||||
|
||||
# if it doesn't exist, create it
|
||||
if context_node is None:
|
||||
context_node = GraphNode(
|
||||
namespace=f"context/{refiners_module.context}/{refiners_module.key}",
|
||||
id=f"context.{refiners_module.context}.{refiners_module.key}",
|
||||
label=refiners_module.key,
|
||||
)
|
||||
|
||||
# connect the nodes: node -> context_node
|
||||
context_node.incomingEdges.append(
|
||||
IncomingEdge(
|
||||
sourceNodeId=node.id,
|
||||
),
|
||||
)
|
||||
graph.nodes.append(context_node)
|
||||
if isinstance(refiners_module, fl.UseContext):
|
||||
# find the context node
|
||||
context_node = find_node(
|
||||
graph=graph,
|
||||
node_id=f"context.{refiners_module.context}.{refiners_module.key}",
|
||||
)
|
||||
|
||||
# if it doesn't exist, create it
|
||||
if context_node is None:
|
||||
context_node = GraphNode(
|
||||
namespace=f"context/{refiners_module.context}/{refiners_module.key}",
|
||||
id=f"context.{refiners_module.context}.{refiners_module.key}",
|
||||
label=refiners_module.key,
|
||||
)
|
||||
graph.nodes.append(context_node)
|
||||
|
||||
# connect the nodes: context_node -> node
|
||||
node.incomingEdges.append(
|
||||
IncomingEdge(
|
||||
sourceNodeId=f"context.{refiners_module.context}.{refiners_module.key}",
|
||||
),
|
||||
)
|
||||
|
||||
return ModelExplorerGraphs(graphs=[graph])
|
||||
44
src/model_explorer_refiners/main.py
Normal file
44
src/model_explorer_refiners/main.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
import refiners.fluxion.layers as fl
|
||||
from model_explorer import Adapter, AdapterMetadata, server
|
||||
from model_explorer.config import ModelExplorerConfig
|
||||
|
||||
from model_explorer_refiners.convert import convert_model
|
||||
|
||||
|
||||
class RefinersAdapter(Adapter):
|
||||
metadata = AdapterMetadata(
|
||||
id="refiners",
|
||||
name="Refiners adapter",
|
||||
description="Refiners adapter for Model Explorer",
|
||||
source_repo="https://github.com/finegrain/model-explorer-refiners",
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@staticmethod
|
||||
def visualize(
|
||||
model: fl.Chain,
|
||||
host: str = "localhost",
|
||||
port: int = 8080,
|
||||
) -> None:
|
||||
"""Visualize a refiners model in Model Explorer.
|
||||
|
||||
Args:
|
||||
model: the refiners model to visualize.
|
||||
host: the host of the server.
|
||||
port: the port of the server.
|
||||
"""
|
||||
# construct config
|
||||
config = ModelExplorerConfig()
|
||||
graph = convert_model(model)
|
||||
config.graphs_list.append(graph)
|
||||
index = len(config.graphs_list) - 1
|
||||
config.model_sources.append({"url": f"graphs://refiners/{index}"})
|
||||
|
||||
# start server
|
||||
server.start(
|
||||
config=config,
|
||||
host=host,
|
||||
port=port,
|
||||
)
|
||||
17
src/model_explorer_refiners/utils.py
Normal file
17
src/model_explorer_refiners/utils.py
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
from model_explorer.graph_builder import Graph, GraphNode
|
||||
|
||||
|
||||
def find_node(graph: Graph, node_id: str) -> GraphNode | None:
|
||||
"""Find a node in a graph.
|
||||
|
||||
Args:
|
||||
graph: the graph to search.
|
||||
node_id: the id of the node to find.
|
||||
|
||||
Returns:
|
||||
The node node which matches the node_id, if found, otherwise None.
|
||||
"""
|
||||
return next(
|
||||
(context_node for context_node in graph.nodes if context_node.id == node_id),
|
||||
None,
|
||||
)
|
||||
Loading…
Reference in a new issue