First draft of "Find interesting heads" page

hexviz/pages/Identify_interesting_Heads.py

@@ -0,0 +1,35 @@
+import streamlit as st
+
+from hexviz.attention import get_attention, get_sequence, get_structure
+from hexviz.models import Model, ModelType
+from hexviz.plot import plot_tiled_heatmap
+
+st.set_page_config(layout="wide")
+
+
+models = [
+    Model(name=ModelType.TAPE_BERT, layers=12, heads=12),
+    Model(name=ModelType.ZymCTRL, layers=36, heads=16),
+]
+
+selected_model_name = st.sidebar.selectbox("Select a model", [model.name.value for model in models], index=0)
+selected_model = next((model for model in models if model.name.value == selected_model_name), None)
+
+pdb_id = st.sidebar.text_input(
+        label="PDB ID",
+        value="1AKE",
+    )
+
+structure = get_structure(pdb_id)
+chains = list(structure.get_chains())
+
+sequence = get_sequence(chains[0])
+l = len(sequence)
+n_residues = st.sidebar.number_input(f"Residue count (1-{l})",value=50, min_value=1, max_value=l)
+truncated_sequence = sequence[:n_residues]
+
+attention = get_attention(sequence=truncated_sequence, model_type=selected_model.name)
+
+st.subheader("Find interesting heads and layers")
+fig = plot_tiled_heatmap(attention, layer_count=selected_model.layers, head_count=selected_model.heads)
+st.pyplot(fig)

hexviz/plot.py

@@ -0,0 +1,23 @@
+import matplotlib.pyplot as plt
+
+
+def plot_tiled_heatmap(tensor, layer_count=12, head_count=12):
+    tensor = tensor[:layer_count:2, :head_count:2, :, :]  # Slice the tensor according to the provided arguments
+    num_layers = layer_count // 2
+    num_heads = head_count // 2
+    fig, axes = plt.subplots(num_layers, num_heads, figsize=(12, 12))
+    for i in range(num_layers):
+        for j in range(num_heads):
+            axes[i, j].imshow(tensor[i, j].detach().numpy(), cmap='viridis', aspect='auto')
+            axes[i, j].axis('off')
+
+            # Enumerate the axes
+            if i == 0:
+                axes[i, j].set_title(f'Head {j * 2 + 1}', fontsize=10, y=1.05)
+
+    # Add layer labels on the right Y-axis
+    for i in range(num_layers):
+        fig.text(0.98, (num_layers - i - 1) / num_layers + 0.025, f'Layer {i * 2 + 1}', fontsize=10, rotation=0, ha='right', va='center')
+
+    plt.subplots_adjust(wspace=0.1, hspace=0.1)
+    return fig

poetry.lock

@@ -293,6 +293,32 @@ traitlets = ">=5.3"
 [package.extras]
 test = ["pytest"]
 
+[[package]]
+name = "contourpy"
+version = "1.0.7"
+description = "Python library for calculating contours of 2D quadrilateral grids"
+category = "main"
+optional = false
+python-versions = ">=3.8"
+
+[package.dependencies]
+numpy = ">=1.16"
+
+[package.extras]
+bokeh = ["bokeh", "chromedriver", "selenium"]
+docs = ["furo", "sphinx-copybutton"]
+mypy = ["contourpy", "docutils-stubs", "mypy (==0.991)", "types-pillow"]
+test = ["matplotlib", "pillow", "pytest"]
+test-no-images = ["pytest"]
+
+[[package]]
+name = "cycler"
+version = "0.11.0"
+description = "Composable style cycles"
+category = "main"
+optional = false
+python-versions = ">=3.6"
+
 [[package]]
 name = "debugpy"
 version = "1.6.6"
@@ -381,6 +407,28 @@ python-versions = ">=3.7"
 docs = ["furo (>=2022.12.7)", "sphinx-autodoc-typehints (>=1.22,!=1.23.4)", "sphinx (>=6.1.3)"]
 testing = ["covdefaults (>=2.3)", "coverage (>=7.2.1)", "pytest-cov (>=4)", "pytest-timeout (>=2.1)", "pytest (>=7.2.2)"]
 
+[[package]]
+name = "fonttools"
+version = "4.39.3"
+description = "Tools to manipulate font files"
+category = "main"
+optional = false
+python-versions = ">=3.8"
+
+[package.extras]
+all = ["fs (>=2.2.0,<3)", "lxml (>=4.0,<5)", "zopfli (>=0.1.4)", "lz4 (>=1.7.4.2)", "matplotlib", "sympy", "skia-pathops (>=0.5.0)", "uharfbuzz (>=0.23.0)", "brotlicffi (>=0.8.0)", "scipy", "brotli (>=1.0.1)", "munkres", "unicodedata2 (>=15.0.0)", "xattr"]
+graphite = ["lz4 (>=1.7.4.2)"]
+interpolatable = ["scipy", "munkres"]
+lxml = ["lxml (>=4.0,<5)"]
+pathops = ["skia-pathops (>=0.5.0)"]
+plot = ["matplotlib"]
+repacker = ["uharfbuzz (>=0.23.0)"]
+symfont = ["sympy"]
+type1 = ["xattr"]
+ufo = ["fs (>=2.2.0,<3)"]
+unicode = ["unicodedata2 (>=15.0.0)"]
+woff = ["zopfli (>=0.1.4)", "brotlicffi (>=0.8.0)", "brotli (>=1.0.1)"]
+
 [[package]]
 name = "fqdn"
 version = "1.5.1"
@@ -810,6 +858,14 @@ category = "main"
 optional = false
 python-versions = ">=3.7"
 
+[[package]]
+name = "kiwisolver"
+version = "1.4.4"
+description = "A fast implementation of the Cassowary constraint solver"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+
 [[package]]
 name = "lit"
 version = "15.0.7"
@@ -855,6 +911,27 @@ category = "main"
 optional = false
 python-versions = ">=3.7"
 
+[[package]]
+name = "matplotlib"
+version = "3.7.1"
+description = "Python plotting package"
+category = "main"
+optional = false
+python-versions = ">=3.8"
+
+[package.dependencies]
+contourpy = ">=1.0.1"
+cycler = ">=0.10"
+fonttools = ">=4.22.0"
+importlib-resources = {version = ">=3.2.0", markers = "python_version < \"3.10\""}
+kiwisolver = ">=1.0.1"
+numpy = ">=1.20"
+packaging = ">=20.0"
+pillow = ">=6.2.0"
+pyparsing = ">=2.3.1"
+python-dateutil = ">=2.7"
+setuptools_scm = ">=7"
+
 [[package]]
 name = "matplotlib-inline"
 version = "0.1.6"
@@ -1400,6 +1477,17 @@ category = "main"
 optional = false
 python-versions = ">=3.6"
 
+[[package]]
+name = "pyparsing"
+version = "3.0.9"
+description = "pyparsing module - Classes and methods to define and execute parsing grammars"
+category = "main"
+optional = false
+python-versions = ">=3.6.8"
+
+[package.extras]
+diagrams = ["railroad-diagrams", "jinja2"]
+
 [[package]]
 name = "pyrsistent"
 version = "0.19.3"
@@ -1622,6 +1710,23 @@ category = "main"
 optional = false
 python-versions = "*"
 
+[[package]]
+name = "setuptools-scm"
+version = "7.1.0"
+description = "the blessed package to manage your versions by scm tags"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+
+[package.dependencies]
+packaging = ">=20.0"
+tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""}
+typing-extensions = "*"
+
+[package.extras]
+test = ["pytest (>=6.2)", "virtualenv (>20)"]
+toml = ["setuptools (>=42)"]
+
 [[package]]
 name = "six"
 version = "1.16.0"
@@ -1820,7 +1925,7 @@ python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*"
 name = "tomli"
 version = "2.0.1"
 description = "A lil' TOML parser"
-category = "dev"
+category = "main"
 optional = false
 python-versions = ">=3.7"
 
@@ -2134,7 +2239,7 @@ testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "flake8 (<5)", "pytest-co
 [metadata]
 lock-version = "1.1"
 python-versions = ">=3.8.9,<3.9.7 || >3.9.7,<4.0.0"
-content-hash = "f7215288fc64984519d9a9ca7b7db5663b34c7727322547d6c8e9e87fdd1161b"
+content-hash = "4583c3e7b0104de5e04c6438baa75cacda39164b118e28c71828ab1aa7a5eb13"
 
 [metadata.files]
 altair = []
@@ -2189,6 +2294,8 @@ click = []
 cmake = []
 colorama = []
 comm = []
+contourpy = []
+cycler = []
 debugpy = []
 decorator = [
     {file = "decorator-5.1.1-py3-none-any.whl", hash = "sha256:b8c3f85900b9dc423225913c5aace94729fe1fa9763b38939a95226f02d37186"},
@@ -2207,6 +2314,7 @@ exceptiongroup = []
 executing = []
 fastjsonschema = []
 filelock = []
+fonttools = []
 fqdn = []
 gitdb = []
 gitpython = []
@@ -2237,10 +2345,12 @@ jupyter-server = []
 jupyter-server-terminals = []
 jupyterlab-pygments = []
 jupyterlab-widgets = []
+kiwisolver = []
 lit = []
 lmdb = []
 markdown-it-py = []
 markupsafe = []
+matplotlib = []
 matplotlib-inline = []
 mdurl = []
 mistune = []
@@ -2311,6 +2421,7 @@ pycparser = [
 pydeck = []
 pygments = []
 pympler = []
+pyparsing = []
 pyrsistent = []
 pytest = []
 python-dateutil = [
@@ -2371,6 +2482,7 @@ send2trash = [
     {file = "Send2Trash-1.8.0.tar.gz", hash = "sha256:d2c24762fd3759860a0aff155e45871447ea58d2be6bdd39b5c8f966a0c99c2d"},
 ]
 sentencepiece = []
+setuptools-scm = []
 six = [
     {file = "six-1.16.0-py2.py3-none-any.whl", hash = "sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254"},
     {file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"},

pyproject.toml

@@ -13,6 +13,7 @@ transformers = "^4.27.1"
 torch = "^2.0.0"
 sentencepiece = "^0.1.97"
 tape-proteins = "^0.5"
+matplotlib = "^3.7.1"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.2.2"

requirements.txt

@@ -22,6 +22,8 @@ click==8.1.3; python_version >= "3.7" and python_full_version < "3.9.7" or pytho
 cmake==3.26.0; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version >= "3.8.0"
 colorama==0.4.6; python_version >= "3.8" and python_full_version < "3.9.7" and platform_system == "Windows" and python_full_version >= "3.7.0" and sys_platform == "win32" or python_full_version > "3.9.7" and python_version >= "3.8" and platform_system == "Windows" and sys_platform == "win32"
 comm==0.1.2; python_version >= "3.8"
+contourpy==1.0.7; python_version >= "3.8"
+cycler==0.11.0; python_version >= "3.8"
 debugpy==1.6.6; python_version >= "3.8"
 decorator==5.1.1; python_version >= "3.8" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.8"
 defusedxml==0.7.1; python_version >= "3.7" and python_full_version < "3.0.0" or python_full_version >= "3.5.0" and python_version >= "3.7"
@@ -30,13 +32,14 @@ entrypoints==0.4; python_version >= "3.7" and python_full_version < "3.9.7" or p
 executing==1.2.0; python_version >= "3.8"
 fastjsonschema==2.16.3; python_version >= "3.7"
 filelock==3.10.0; python_version >= "3.7" and python_full_version >= "3.8.0" and platform_system == "Linux" and platform_machine == "x86_64"
+fonttools==4.39.3; python_version >= "3.8"
 fqdn==1.5.1; python_version >= "3.7" and python_version < "4" and python_full_version < "3.9.7" or python_version >= "3.7" and python_version < "4" and python_full_version > "3.9.7"
 gitdb==4.0.10; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7"
 gitpython==3.1.31; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7"
 huggingface-hub==0.13.2; python_full_version >= "3.7.0"
 idna==3.4; python_version >= "3.8" and python_version < "4" and (python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7") and python_full_version >= "3.7.0" and (python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7")
 importlib-metadata==6.0.0; python_version >= "3.8" and python_full_version < "3.9.7" and python_version < "3.10" or python_full_version > "3.9.7" and python_version >= "3.8" and python_version < "3.10"
-importlib-resources==5.12.0; python_version >= "3.7" and python_full_version < "3.9.7" and python_version < "3.9" or python_full_version > "3.9.7" and python_version >= "3.7" and python_version < "3.9"
+importlib-resources==5.12.0; python_version < "3.10" and python_version >= "3.8" and (python_version >= "3.7" and python_full_version < "3.9.7" and python_version < "3.9" or python_full_version > "3.9.7" and python_version >= "3.7" and python_version < "3.9")
 ipykernel==6.21.3; python_version >= "3.8"
 ipyspeck==0.6.1; python_version >= "3.6"
 ipython-genutils==0.2.0; python_version >= "3.7"
@@ -56,11 +59,13 @@ jupyter-server-terminals==0.4.4; python_version >= "3.8"
 jupyter-server==2.4.0; python_version >= "3.8"
 jupyterlab-pygments==0.2.2; python_version >= "3.7"
 jupyterlab-widgets==3.0.5; python_version >= "3.7"
+kiwisolver==1.4.4; python_version >= "3.8"
 lit==15.0.7; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version >= "3.8.0"
 lmdb==1.4.0
 markdown-it-py==2.2.0; python_version >= "3.7" and python_full_version < "3.9.7" and python_full_version >= "3.7.0" or python_full_version > "3.9.7" and python_version >= "3.7"
 markupsafe==2.1.2; python_version >= "3.7" and python_full_version >= "3.8.0" and (python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7")
 matplotlib-inline==0.1.6; python_version >= "3.8"
+matplotlib==3.7.1; python_version >= "3.8"
 mdurl==0.1.2; python_version >= "3.7" and python_full_version < "3.9.7" and python_full_version >= "3.7.0" or python_full_version > "3.9.7" and python_version >= "3.7"
 mistune==2.0.5; python_version >= "3.7"
 mpmath==1.3.0; python_version >= "3.8" and python_full_version >= "3.8.0"
@@ -90,7 +95,7 @@ pandocfilters==1.5.0; python_version >= "3.7" and python_full_version < "3.0.0"
 parso==0.8.3; python_version >= "3.8"
 pexpect==4.8.0; sys_platform != "win32" and python_version >= "3.8"
 pickleshare==0.7.5; python_version >= "3.8"
-pillow==9.4.0; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7"
+pillow==9.4.0; python_version >= "3.8" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.8"
 pkgutil-resolve-name==1.3.10; python_version >= "3.7" and python_full_version < "3.9.7" and python_version < "3.9" or python_full_version > "3.9.7" and python_version >= "3.7" and python_version < "3.9"
 platformdirs==3.1.1; python_version >= "3.8"
 prometheus-client==0.16.0; python_version >= "3.8"
@@ -105,6 +110,7 @@ pycparser==2.21; python_version >= "3.8" and python_full_version < "3.0.0" and i
 pydeck==0.8.0; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7"
 pygments==2.14.0; python_version >= "3.8" and python_full_version < "3.9.7" and python_full_version >= "3.7.0" or python_full_version > "3.9.7" and python_version >= "3.8"
 pympler==1.0.1; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.6"
+pyparsing==3.0.9; python_full_version >= "3.6.8" and python_version >= "3.8"
 pyrsistent==0.19.3; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7"
 python-dateutil==2.8.2; python_version >= "3.8" and python_full_version < "3.0.0" or python_version >= "3.8" and python_full_version < "3.9.7" and python_full_version >= "3.3.0" or python_full_version > "3.9.7" and python_version >= "3.8"
 python-json-logger==2.0.7; python_version >= "3.8"
@@ -124,6 +130,7 @@ scipy==1.9.3; python_version >= "3.8"
 semver==2.13.0; python_version >= "3.7" and python_full_version < "3.0.0" or python_version >= "3.7" and python_full_version < "3.9.7" and python_full_version >= "3.4.0" or python_full_version > "3.9.7" and python_version >= "3.6"
 send2trash==1.8.0; python_version >= "3.8"
 sentencepiece==0.1.97
+setuptools-scm==7.1.0; python_version >= "3.8"
 six==1.16.0; python_version >= "3.8" and python_full_version < "3.0.0" or python_version >= "3.8" and python_full_version < "3.9.7" and python_full_version >= "3.5.0" or python_full_version > "3.9.7" and python_version >= "3.8"
 smmap==5.0.0; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7"
 sniffio==1.3.0; python_full_version >= "3.6.2" and python_version >= "3.8"
@@ -138,6 +145,7 @@ terminado==0.17.1; python_version >= "3.8"
 tinycss2==1.2.1; python_version >= "3.7"
 tokenizers==0.13.2; python_full_version >= "3.7.0"
 toml==0.10.2; python_version >= "3.7" and python_full_version < "3.0.0" or python_version >= "3.7" and python_full_version < "3.9.7" and python_full_version >= "3.3.0" or python_full_version > "3.9.7" and python_version >= "3.6"
+tomli==2.0.1; python_version < "3.11" and python_version >= "3.8"
 tomlkit==0.11.6; python_version >= "3.7"
 toolz==0.12.0; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7"
 torch==2.0.0; python_full_version >= "3.8.0"
@@ -146,7 +154,7 @@ tqdm==4.65.0; python_version >= "3.7" and python_full_version >= "3.7.0"
 traitlets==5.9.0; python_full_version >= "3.7.0" and python_version >= "3.8"
 transformers==4.27.1; python_full_version >= "3.7.0"
 triton==2.0.0; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version >= "3.8.0"
-typing-extensions==4.5.0; python_version >= "3.7" and python_full_version < "3.9.7" and python_full_version >= "3.8.0" and python_version < "3.9" and platform_system == "Linux" and platform_machine == "x86_64" or python_full_version > "3.9.7" and python_version >= "3.7" and python_version < "3.9" and platform_system == "Linux" and platform_machine == "x86_64"
+typing-extensions==4.5.0; python_version >= "3.8" and python_full_version < "3.9.7" and python_full_version >= "3.8.0" and python_version < "3.9" and platform_system == "Linux" and platform_machine == "x86_64" or python_full_version > "3.9.7" and python_version >= "3.8" and python_version < "3.9" and platform_system == "Linux" and platform_machine == "x86_64"
 tzdata==2022.7; python_version >= "3.7" and python_full_version < "3.0.0" and platform_system == "Windows" or python_version >= "3.7" and python_full_version < "3.9.7" and platform_system == "Windows" and python_full_version >= "3.6.0" or python_full_version > "3.9.7" and python_version >= "3.6" and platform_system == "Windows"
 tzlocal==4.2; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.6"
 uri-template==1.2.0; python_version >= "3.7" and python_full_version < "3.9.7" or python_full_version > "3.9.7" and python_version >= "3.7"
@@ -158,4 +166,4 @@ webcolors==1.12; python_version >= "3.7" and python_full_version < "3.9.7" or py
 webencodings==0.5.1; python_version >= "3.7"
 websocket-client==1.5.1; python_version >= "3.8"
 widgetsnbextension==3.5.2; python_version >= "3.6"
-zipp==3.15.0; python_version >= "3.7" and python_full_version < "3.9.7" and python_version < "3.9" or python_full_version > "3.9.7" and python_version >= "3.7" and python_version < "3.9"
+zipp==3.15.0; python_version >= "3.8" and python_full_version < "3.9.7" and python_version < "3.10" or python_full_version > "3.9.7" and python_version >= "3.8" and python_version < "3.10"