Split attention.py into attention.py and models.py

hexviz/app.py

@@ -3,7 +3,8 @@ import stmol
 import streamlit as st
 from stmol import showmol
 
-from hexviz.attention import Model, ModelType, get_attention_pairs
+from hexviz.attention import get_attention_pairs
+from hexviz.models import Model, ModelType
 
 st.title("Attention Visualization on proteins")
 
@@ -12,8 +13,6 @@ Visualize attention weights on protein structures for the protein language model
 Pick a PDB ID, layer and head to visualize attention.
 """
 
-
-# Define list of model types
 models = [
     # Model(name=ModelType.ProtGPT2, layers=36, heads=20),
     Model(name=ModelType.TAPE_BERT, layers=12, heads=12),

hexviz/attention.py

@@ -6,24 +6,10 @@ from urllib import request
 import streamlit as st
 import torch
 from Bio.PDB import PDBParser, Polypeptide, Structure
-from tape import ProteinBertModel, TAPETokenizer
-from transformers import (AutoTokenizer, GPT2LMHeadModel, T5EncoderModel,
-                          T5Tokenizer)
+from models import (ModelType, get_protgpt2, get_protT5, get_tape_bert,
+                    get_zymctrl)
 
 
-class ModelType(str, Enum):
-    TAPE_BERT = "TAPE-BERT"
-    PROT_T5 = "prot_t5_xl_half_uniref50-enc"
-    ZymCTRL = "ZymCTRL"
-    ProtGPT2 = "ProtGPT2"
-
-
-class Model:
-    def __init__(self, name, layers, heads):
-        self.name: ModelType = name
-        self.layers: int = layers
-        self.heads: int = heads
-
 @st.cache
 def get_structure(pdb_code: str) -> Structure:
     """
@@ -51,41 +37,6 @@ def get_sequences(structure: Structure) -> List[str]:
         sequences.append("".join(list(residues_single_letter)))
     return sequences
 
-@st.cache
-def get_protT5() -> Tuple[T5Tokenizer, T5EncoderModel]:
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    tokenizer = T5Tokenizer.from_pretrained(
-        "Rostlab/prot_t5_xl_half_uniref50-enc", do_lower_case=False
-    )
-
-    model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc").to(
-        device
-    )
-
-    model.full() if device == "cpu" else model.half()
-
-    return tokenizer, model
-
-@st.cache
-def get_tape_bert() -> Tuple[TAPETokenizer, ProteinBertModel]:
-    tokenizer = TAPETokenizer()
-    model = ProteinBertModel.from_pretrained('bert-base', output_attentions=True)
-    return tokenizer, model
-
-@st.cache
-def get_zymctrl() -> Tuple[AutoTokenizer, GPT2LMHeadModel]:
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    tokenizer = AutoTokenizer.from_pretrained('nferruz/ZymCTRL')
-    model = GPT2LMHeadModel.from_pretrained('nferruz/ZymCTRL').to(device)
-    return tokenizer, model
-
-@st.cache
-def get_protgpt2() -> Tuple[AutoTokenizer, GPT2LMHeadModel]:
-    device = torch.device('cuda')
-    tokenizer = AutoTokenizer.from_pretrained('nferruz/ProtGPT2')
-    model = GPT2LMHeadModel.from_pretrained('nferruz/ProtGPT2').to(device)
-    return tokenizer, model
-
 @st.cache
 def get_attention(
     sequence: str, model_type: ModelType = ModelType.TAPE_BERT  

hexviz/models.py

@@ -0,0 +1,57 @@
+from enum import Enum
+from typing import Tuple
+
+import streamlit as st
+import torch
+from tape import ProteinBertModel, TAPETokenizer
+from transformers import (AutoTokenizer, GPT2LMHeadModel, T5EncoderModel,
+                          T5Tokenizer)
+
+
+class ModelType(str, Enum):
+    TAPE_BERT = "TAPE-BERT"
+    PROT_T5 = "prot_t5_xl_half_uniref50-enc"
+    ZymCTRL = "ZymCTRL"
+    ProtGPT2 = "ProtGPT2"
+
+
+class Model:
+    def __init__(self, name, layers, heads):
+        self.name: ModelType = name
+        self.layers: int = layers
+        self.heads: int = heads
+
+@st.cache
+def get_protT5() -> Tuple[T5Tokenizer, T5EncoderModel]:
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    tokenizer = T5Tokenizer.from_pretrained(
+        "Rostlab/prot_t5_xl_half_uniref50-enc", do_lower_case=False
+    )
+
+    model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc").to(
+        device
+    )
+
+    model.full() if device == "cpu" else model.half()
+
+    return tokenizer, model
+
+@st.cache
+def get_tape_bert() -> Tuple[TAPETokenizer, ProteinBertModel]:
+    tokenizer = TAPETokenizer()
+    model = ProteinBertModel.from_pretrained('bert-base', output_attentions=True)
+    return tokenizer, model
+
+@st.cache
+def get_zymctrl() -> Tuple[AutoTokenizer, GPT2LMHeadModel]:
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained('nferruz/ZymCTRL')
+    model = GPT2LMHeadModel.from_pretrained('nferruz/ZymCTRL').to(device)
+    return tokenizer, model
+
+@st.cache
+def get_protgpt2() -> Tuple[AutoTokenizer, GPT2LMHeadModel]:
+    device = torch.device('cuda')
+    tokenizer = AutoTokenizer.from_pretrained('nferruz/ProtGPT2')
+    model = GPT2LMHeadModel.from_pretrained('nferruz/ProtGPT2').to(device)
+    return tokenizer, model

tests/test_attention.py

@@ -1,11 +1,8 @@
 import torch
 from Bio.PDB.Structure import Structure
-from transformers import (GPT2LMHeadModel, GPT2TokenizerFast, T5EncoderModel,
-                          T5Tokenizer)
 
-from hexviz.attention import (ModelType, get_attention, get_protT5,
-                              get_sequences, get_structure, get_zymctrl,
-                              unidirectional_sum_filtered)
+from hexviz.attention import (ModelType, get_attention, get_sequences,
+                              get_structure, unidirectional_sum_filtered)
 
 
 def test_get_structure():
@@ -27,27 +24,6 @@ def test_get_sequences():
     A, B = sequences
     assert A[:3] == ["M", "R", "I"]
 
-def test_get_protT5():
-    result = get_protT5()
-
-    assert result is not None
-    assert isinstance(result, tuple)
-
-    tokenizer, model = result
-
-    assert isinstance(tokenizer, T5Tokenizer)
-    assert isinstance(model, T5EncoderModel)
-
-def test_get_zymctrl():
-    result = get_zymctrl()
-
-    assert result is not None
-    assert isinstance(result, tuple)
-
-    tokenizer, model = result
-
-    assert isinstance(tokenizer, GPT2TokenizerFast)
-    assert isinstance(model, GPT2LMHeadModel)
 
 def test_get_attention_zymctrl():
 

tests/test_models.py

@@ -0,0 +1,28 @@
+
+from transformers import (GPT2LMHeadModel, GPT2TokenizerFast, T5EncoderModel,
+                          T5Tokenizer)
+
+from hexviz.models import get_protT5, get_zymctrl
+
+
+def test_get_protT5():
+    result = get_protT5()
+
+    assert result is not None
+    assert isinstance(result, tuple)
+
+    tokenizer, model = result
+
+    assert isinstance(tokenizer, T5Tokenizer)
+    assert isinstance(model, T5EncoderModel)
+
+def test_get_zymctrl():
+    result = get_zymctrl()
+
+    assert result is not None
+    assert isinstance(result, tuple)
+
+    tokenizer, model = result
+
+    assert isinstance(tokenizer, GPT2TokenizerFast)
+    assert isinstance(model, GPT2LMHeadModel)