Upload 26 files

tape/__init__.py

@@ -0,0 +1,21 @@
+from . import datasets  # noqa: F401
+from . import metrics  # noqa: F401
+from .tokenizers import TAPETokenizer  # noqa: F401
+from .models.modeling_utils import ProteinModel
+from .models.modeling_utils import ProteinConfig
+
+import sys
+from pathlib import Path
+import importlib
+import pkgutil
+
+__version__ = '0.4'
+
+
+# Import all the models and configs
+for _, name, _ in pkgutil.iter_modules([str(Path(__file__).parent / 'models')]):
+    imported_module = importlib.import_module('.models.' + name, package=__name__)
+    for name, cls in imported_module.__dict__.items():
+        if isinstance(cls, type) and \
+                (issubclass(cls, ProteinModel) or issubclass(cls, ProteinConfig)):
+            setattr(sys.modules[__name__], name, cls)

tape/datasets.py

@@ -0,0 +1,926 @@
+from typing import Union, List, Tuple, Sequence, Dict, Any, Optional, Collection
+from copy import copy
+from pathlib import Path
+import pickle as pkl
+import logging
+import random
+
+import lmdb
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from scipy.spatial.distance import pdist, squareform
+
+from .tokenizers import TAPETokenizer
+from .registry import registry
+
+logger = logging.getLogger(__name__)
+
+
+def dataset_factory(data_file: Union[str, Path], *args, **kwargs) -> Dataset:
+    data_file = Path(data_file)
+    if not data_file.exists():
+        raise FileNotFoundError(data_file)
+    if data_file.suffix == '.lmdb':
+        return LMDBDataset(data_file, *args, **kwargs)
+    elif data_file.suffix in {'.fasta', '.fna', '.ffn', '.faa', '.frn'}:
+        return FastaDataset(data_file, *args, **kwargs)
+    elif data_file.suffix == '.json':
+        return JSONDataset(data_file, *args, **kwargs)
+    elif data_file.is_dir():
+        return NPZDataset(data_file, *args, **kwargs)
+    else:
+        raise ValueError(f"Unrecognized datafile type {data_file.suffix}")
+
+
+def pad_sequences(sequences: Sequence, constant_value=0, dtype=None) -> np.ndarray:
+    batch_size = len(sequences)
+    shape = [batch_size] + np.max([seq.shape for seq in sequences], 0).tolist()
+
+    if dtype is None:
+        dtype = sequences[0].dtype
+
+    if isinstance(sequences[0], np.ndarray):
+        array = np.full(shape, constant_value, dtype=dtype)
+    elif isinstance(sequences[0], torch.Tensor):
+        array = torch.full(shape, constant_value, dtype=dtype)
+
+    for arr, seq in zip(array, sequences):
+        arrslice = tuple(slice(dim) for dim in seq.shape)
+        arr[arrslice] = seq
+
+    return array
+
+
+class FastaDataset(Dataset):
+    """Creates a dataset from a fasta file.
+    Args:
+        data_file (Union[str, Path]): Path to fasta file.
+        in_memory (bool, optional): Whether to load the full dataset into memory.
+            Default: False.
+    """
+
+    def __init__(self,
+                 data_file: Union[str, Path],
+                 in_memory: bool = False):
+
+        from Bio import SeqIO
+        data_file = Path(data_file)
+        if not data_file.exists():
+            raise FileNotFoundError(data_file)
+
+        # if in_memory:
+        cache = list(SeqIO.parse(str(data_file), 'fasta'))
+        num_examples = len(cache)
+        self._cache = cache
+        # else:
+            # records = SeqIO.index(str(data_file), 'fasta')
+            # num_examples = len(records)
+#
+            # if num_examples < 10000:
+                # logger.info("Reading full fasta file into memory because number of examples "
+                            # "is very low. This loads data approximately 20x faster.")
+                # in_memory = True
+                # cache = list(records.values())
+                # self._cache = cache
+            # else:
+                # self._records = records
+                # self._keys = list(records.keys())
+
+        self._in_memory = in_memory
+        self._num_examples = num_examples
+
+    def __len__(self) -> int:
+        return self._num_examples
+
+    def __getitem__(self, index: int):
+        if not 0 <= index < self._num_examples:
+            raise IndexError(index)
+
+        # if self._in_memory and self._cache[index] is not None:
+        record = self._cache[index]
+        # else:
+            # key = self._keys[index]
+            # record = self._records[key]
+            # if self._in_memory:
+                # self._cache[index] = record
+
+        item = {'id': record.id,
+                'primary': str(record.seq),
+                'protein_length': len(record.seq)}
+        return item
+
+
+class LMDBDataset(Dataset):
+    """Creates a dataset from an lmdb file.
+    Args:
+        data_file (Union[str, Path]): Path to lmdb file.
+        in_memory (bool, optional): Whether to load the full dataset into memory.
+            Default: False.
+    """
+
+    def __init__(self,
+                 data_file: Union[str, Path],
+                 in_memory: bool = False):
+
+        data_file = Path(data_file)
+        if not data_file.exists():
+            raise FileNotFoundError(data_file)
+
+        env = lmdb.open(str(data_file), max_readers=1, readonly=True,
+                        lock=False, readahead=False, meminit=False)
+
+        with env.begin(write=False) as txn:
+            num_examples = pkl.loads(txn.get(b'num_examples'))
+
+        if in_memory:
+            cache = [None] * num_examples
+            self._cache = cache
+
+        self._env = env
+        self._in_memory = in_memory
+        self._num_examples = num_examples
+
+    def __len__(self) -> int:
+        return self._num_examples
+
+    def __getitem__(self, index: int):
+        if not 0 <= index < self._num_examples:
+            raise IndexError(index)
+
+        if self._in_memory and self._cache[index] is not None:
+            item = self._cache[index]
+        else:
+            with self._env.begin(write=False) as txn:
+                item = pkl.loads(txn.get(str(index).encode()))
+                if 'id' not in item:
+                    item['id'] = str(index)
+                if self._in_memory:
+                    self._cache[index] = item
+        return item
+
+
+class JSONDataset(Dataset):
+    """Creates a dataset from a json file. Assumes that data is
+       a JSON serialized list of record, where each record is
+       a dictionary.
+    Args:
+        data_file (Union[str, Path]): Path to json file.
+        in_memory (bool): Dummy variable to match API of other datasets
+    """
+
+    def __init__(self, data_file: Union[str, Path], in_memory: bool = True):
+        import json
+        data_file = Path(data_file)
+        if not data_file.exists():
+            raise FileNotFoundError(data_file)
+        records = json.loads(data_file.read_text())
+
+        if not isinstance(records, list):
+            raise TypeError(f"TAPE JSONDataset requires a json serialized list, "
+                            f"received {type(records)}")
+        self._records = records
+        self._num_examples = len(records)
+
+    def __len__(self) -> int:
+        return self._num_examples
+
+    def __getitem__(self, index: int):
+        if not 0 <= index < self._num_examples:
+            raise IndexError(index)
+
+        item = self._records[index]
+        if not isinstance(item, dict):
+            raise TypeError(f"Expected dataset to contain a list of dictionary "
+                            f"records, received record of type {type(item)}")
+        if 'id' not in item:
+            item['id'] = str(index)
+        return item
+
+
+class NPZDataset(Dataset):
+    """Creates a dataset from a directory of npz files.
+    Args:
+        data_file (Union[str, Path]): Path to directory of npz files
+        in_memory (bool): Dummy variable to match API of other datasets
+    """
+
+    def __init__(self,
+                 data_file: Union[str, Path],
+                 in_memory: bool = True,
+                 split_files: Optional[Collection[str]] = None):
+        data_file = Path(data_file)
+        if not data_file.exists():
+            raise FileNotFoundError(data_file)
+        if not data_file.is_dir():
+            raise NotADirectoryError(data_file)
+        file_glob = data_file.glob('*.npz')
+        if split_files is None:
+            file_list = list(file_glob)
+        else:
+            split_files = set(split_files)
+            if len(split_files) == 0:
+                raise ValueError("Passed an empty split file set")
+
+            file_list = [f for f in file_glob if f.name in split_files]
+            if len(file_list) != len(split_files):
+                num_missing = len(split_files) - len(file_list)
+                raise FileNotFoundError(
+                    f"{num_missing} specified split files not found in directory")
+
+        if len(file_list) == 0:
+            raise FileNotFoundError(f"No .npz files found in {data_file}")
+
+        self._file_list = file_list
+
+    def __len__(self) -> int:
+        return len(self._file_list)
+
+    def __getitem__(self, index: int):
+        if not 0 <= index < len(self):
+            raise IndexError(index)
+
+        item = dict(np.load(self._file_list[index]))
+        if not isinstance(item, dict):
+            raise TypeError(f"Expected dataset to contain a list of dictionary "
+                            f"records, received record of type {type(item)}")
+        if 'id' not in item:
+            item['id'] = self._file_list[index].stem
+        return item
+
+
+@registry.register_task('embed')
+class EmbedDataset(Dataset):
+
+    def __init__(self,
+                 data_file: Union[str, Path],
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False,
+                 convert_tokens_to_ids: bool = True):
+        super().__init__()
+
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+        self.data = dataset_factory(data_file)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index: int):
+        item = self.data[index]
+        token_ids = self.tokenizer.encode(item['primary'])
+        input_mask = np.ones_like(token_ids)
+        return item['id'], token_ids, input_mask
+
+    def collate_fn(self, batch: List[Tuple[Any, ...]]) -> Dict[str, torch.Tensor]:
+        ids, tokens, input_mask = zip(*batch)
+        ids = list(ids)
+        tokens = torch.from_numpy(pad_sequences(tokens))
+        input_mask = torch.from_numpy(pad_sequences(input_mask))
+        return {'ids': ids, 'input_ids': tokens, 'input_mask': input_mask}  # type: ignore
+
+
+@registry.register_task('masked_language_modeling')
+class MaskedLanguageModelingDataset(Dataset):
+    """Creates the Masked Language Modeling Pfam Dataset
+    Args:
+        data_path (Union[str, Path]): Path to tape data root.
+        split (str): One of ['train', 'valid', 'holdout'], specifies which data file to load.
+        in_memory (bool, optional): Whether to load the full dataset into memory.
+            Default: False.
+    """
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+        super().__init__()
+        if split not in ('train', 'valid', 'holdout'):
+            raise ValueError(
+                f"Unrecognized split: {split}. "
+                f"Must be one of ['train', 'valid', 'holdout']")
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+
+        data_path = Path(data_path)
+        data_file = f'pfam/pfam_{split}.lmdb'
+        self.data = dataset_factory(data_path / data_file, in_memory)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index):
+        item = self.data[index]
+        tokens = self.tokenizer.tokenize(item['primary'])
+        tokens = self.tokenizer.add_special_tokens(tokens)
+        masked_tokens, labels = self._apply_bert_mask(tokens)
+        masked_token_ids = np.array(
+            self.tokenizer.convert_tokens_to_ids(masked_tokens), np.int64)
+        input_mask = np.ones_like(masked_token_ids)
+
+        masked_token_ids = np.array(
+            self.tokenizer.convert_tokens_to_ids(masked_tokens), np.int64)
+
+        return masked_token_ids, input_mask, labels, item['clan'], item['family']
+
+    def collate_fn(self, batch: List[Any]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, lm_label_ids, clan, family = tuple(zip(*batch))
+
+        input_ids = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        # ignore_index is -1
+        lm_label_ids = torch.from_numpy(pad_sequences(lm_label_ids, -1))
+        clan = torch.LongTensor(clan)  # type: ignore
+        family = torch.LongTensor(family)  # type: ignore
+
+        return {'input_ids': input_ids,
+                'input_mask': input_mask,
+                'targets': lm_label_ids}
+
+    def _apply_bert_mask(self, tokens: List[str]) -> Tuple[List[str], List[int]]:
+        masked_tokens = copy(tokens)
+        labels = np.zeros([len(tokens)], np.int64) - 1
+
+        for i, token in enumerate(tokens):
+            # Tokens begin and end with start_token and stop_token, ignore these
+            if token in (self.tokenizer.start_token, self.tokenizer.stop_token):
+                pass
+
+            prob = random.random()
+            if prob < 0.15:
+                prob /= 0.15
+                labels[i] = self.tokenizer.convert_token_to_id(token)
+
+                if prob < 0.8:
+                    # 80% random change to mask token
+                    token = self.tokenizer.mask_token
+                elif prob < 0.9:
+                    # 10% chance to change to random token
+                    token = self.tokenizer.convert_id_to_token(
+                        random.randint(0, self.tokenizer.vocab_size - 1))
+                else:
+                    # 10% chance to keep current token
+                    pass
+
+                masked_tokens[i] = token
+
+        return masked_tokens, labels
+
+
+@registry.register_task('beta_lactamase')
+class BetaModelingDataset(MaskedLanguageModelingDataset):
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+        super().__init__(data_path, split, tokenizer, in_memory)
+        data_path = Path(data_path)
+        data_file = f'unilanguage/{split}_combined.fasta'
+        self.data = dataset_factory(data_path / data_file, in_memory)
+
+    def collate_fn(self, batch: List[Any]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, lm_label_ids = tuple(zip(*batch))
+
+        input_ids = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        # ignore_index is -1
+        lm_label_ids = torch.from_numpy(pad_sequences(lm_label_ids, -1))
+
+        return {'input_ids': input_ids,
+                'input_mask': input_mask,
+                'targets': lm_label_ids}
+
+    def __getitem__(self, index):
+        item = self.data[index]
+        tokens = self.tokenizer.tokenize(item['primary'])
+        tokens = self.tokenizer.add_special_tokens(tokens)
+        masked_tokens, labels = self._apply_bert_mask(tokens)
+        masked_token_ids = np.array(
+            self.tokenizer.convert_tokens_to_ids(masked_tokens), np.int64)
+        input_mask = np.ones_like(masked_token_ids)
+
+        masked_token_ids = np.array(
+            self.tokenizer.convert_tokens_to_ids(masked_tokens), np.int64)
+
+        return masked_token_ids, input_mask, labels
+
+
+@registry.register_task('unilanguage')
+class UniModelingDataset(MaskedLanguageModelingDataset):
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+        super().__init__(data_path, split, tokenizer, in_memory)
+        data_path = Path(data_path)
+        data_file = f'unilanguage/PF00144_full_length_sequences_labeled.fasta'
+        self.data = dataset_factory(data_path / data_file, in_memory)    
+
+    def collate_fn(self, batch: List[Any]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, lm_label_ids = tuple(zip(*batch))
+
+        input_ids = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        # ignore_index is -1
+        lm_label_ids = torch.from_numpy(pad_sequences(lm_label_ids, -1))
+
+        return {'input_ids': input_ids,
+                'input_mask': input_mask,
+                'targets': lm_label_ids}
+
+    def __getitem__(self, index):
+        item = self.data[index]
+        tokens = self.tokenizer.tokenize(item['primary'])
+        tokens = self.tokenizer.add_special_tokens(tokens)
+        masked_tokens, labels = self._apply_bert_mask(tokens)
+        masked_token_ids = np.array(
+            self.tokenizer.convert_tokens_to_ids(masked_tokens), np.int64)
+        input_mask = np.ones_like(masked_token_ids)
+
+        masked_token_ids = np.array(
+            self.tokenizer.convert_tokens_to_ids(masked_tokens), np.int64)
+
+        return masked_token_ids, input_mask, labels
+
+
+@registry.register_task('language_modeling')
+class LanguageModelingDataset(Dataset):
+    """Creates the Language Modeling Pfam Dataset
+    Args:
+        data_path (Union[str, Path]): Path to tape data root.
+        split (str): One of ['train', 'valid', 'holdout'], specifies which data file to load.
+        in_memory (bool, optional): Whether to load the full dataset into memory.
+            Default: False.
+    """
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+        super().__init__()
+        if split not in ('train', 'valid', 'holdout'):
+            raise ValueError(
+                f"Unrecognized split: {split}. "
+                f"Must be one of ['train', 'valid', 'holdout']")
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+
+        data_path = Path(data_path)
+        data_file = f'pfam/pfam_{split}.lmdb'
+        self.data = dataset_factory(data_path / data_file, in_memory)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index):
+        item = self.data[index]
+        token_ids = self.tokenizer.encode(item['primary'])
+        input_mask = np.ones_like(token_ids)
+
+        return token_ids, input_mask, item['clan'], item['family']
+
+    def collate_fn(self, batch: List[Any]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, clan, family = tuple(zip(*batch))
+
+        torch_inputs = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        # ignore_index is -1
+        torch_labels = torch.from_numpy(pad_sequences(input_ids, -1))
+        clan = torch.LongTensor(clan)  # type: ignore
+        family = torch.LongTensor(family)  # type: ignore
+
+        return {'input_ids': torch_inputs,
+                'input_mask': input_mask,
+                'targets': torch_labels}
+
+
+@registry.register_task('fluorescence')
+class FluorescenceDataset(Dataset):
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+
+        if split not in ('train', 'valid', 'test'):
+            raise ValueError(f"Unrecognized split: {split}. "
+                             f"Must be one of ['train', 'valid', 'test']")
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+
+        data_path = Path(data_path)
+        data_file = f'fluorescence/fluorescence_{split}.lmdb'
+        self.data = dataset_factory(data_path / data_file, in_memory)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index: int):
+        item = self.data[index]
+        token_ids = self.tokenizer.encode(item['primary'])
+        input_mask = np.ones_like(token_ids)
+        return token_ids, input_mask, float(item['log_fluorescence'][0])
+
+    def collate_fn(self, batch: List[Tuple[Any, ...]]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, fluorescence_true_value = tuple(zip(*batch))
+        input_ids = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        fluorescence_true_value = torch.FloatTensor(fluorescence_true_value)  # type: ignore
+        fluorescence_true_value = fluorescence_true_value.unsqueeze(1)
+
+        return {'input_ids': input_ids,
+                'input_mask': input_mask,
+                'targets': fluorescence_true_value}
+
+
+@registry.register_task('stability')
+class StabilityDataset(Dataset):
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+
+        if split not in ('train', 'valid', 'test'):
+            raise ValueError(f"Unrecognized split: {split}. "
+                             f"Must be one of ['train', 'valid', 'test']")
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+
+        data_path = Path(data_path)
+        data_file = f'stability/stability_{split}.lmdb'
+
+        self.data = dataset_factory(data_path / data_file, in_memory)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index: int):
+        item = self.data[index]
+        token_ids = self.tokenizer.encode(item['primary'])
+        input_mask = np.ones_like(token_ids)
+        return token_ids, input_mask, float(item['stability_score'][0])
+
+    def collate_fn(self, batch: List[Tuple[Any, ...]]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, stability_true_value = tuple(zip(*batch))
+        input_ids = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        stability_true_value = torch.FloatTensor(stability_true_value)  # type: ignore
+        stability_true_value = stability_true_value.unsqueeze(1)
+
+        return {'input_ids': input_ids,
+                'input_mask': input_mask,
+                'targets': stability_true_value}
+
+
+@registry.register_task('remote_homology', num_labels=1195)
+class RemoteHomologyDataset(Dataset):
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+
+        if split not in ('train', 'valid', 'test_fold_holdout',
+                         'test_family_holdout', 'test_superfamily_holdout'):
+            raise ValueError(f"Unrecognized split: {split}. Must be one of "
+                             f"['train', 'valid', 'test_fold_holdout', "
+                             f"'test_family_holdout', 'test_superfamily_holdout']")
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+
+        data_path = Path(data_path)
+        data_file = f'remote_homology/remote_homology_{split}.lmdb'
+        self.data = dataset_factory(data_path / data_file, in_memory)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index: int):
+        item = self.data[index]
+        token_ids = self.tokenizer.encode(item['primary'])
+        input_mask = np.ones_like(token_ids)
+        return token_ids, input_mask, item['fold_label']
+
+    def collate_fn(self, batch: List[Tuple[Any, ...]]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, fold_label = tuple(zip(*batch))
+        input_ids = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        fold_label = torch.LongTensor(fold_label)  # type: ignore
+
+        return {'input_ids': input_ids,
+                'input_mask': input_mask,
+                'targets': fold_label}
+
+
+@registry.register_task('contact_prediction')
+class ProteinnetDataset(Dataset):
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+
+        if split not in ('train', 'train_unfiltered', 'valid', 'test'):
+            raise ValueError(f"Unrecognized split: {split}. Must be one of "
+                             f"['train', 'train_unfiltered', 'valid', 'test']")
+
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+
+        data_path = Path(data_path)
+        data_file = f'proteinnet/proteinnet_{split}.lmdb'
+        self.data = dataset_factory(data_path / data_file, in_memory)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index: int):
+        item = self.data[index]
+        protein_length = len(item['primary'])
+        token_ids = self.tokenizer.encode(item['primary'])
+        input_mask = np.ones_like(token_ids)
+
+        valid_mask = item['valid_mask']
+        contact_map = np.less(squareform(pdist(item['tertiary'])), 8.0).astype(np.int64)
+
+        yind, xind = np.indices(contact_map.shape)
+        invalid_mask = ~(valid_mask[:, None] & valid_mask[None, :])
+        invalid_mask |= np.abs(yind - xind) < 6
+        contact_map[invalid_mask] = -1
+
+        return token_ids, input_mask, contact_map, protein_length
+
+    def collate_fn(self, batch: List[Tuple[Any, ...]]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, contact_labels, protein_length = tuple(zip(*batch))
+        input_ids = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        contact_labels = torch.from_numpy(pad_sequences(contact_labels, -1))
+        protein_length = torch.LongTensor(protein_length)  # type: ignore
+
+        return {'input_ids': input_ids,
+                'input_mask': input_mask,
+                'targets': contact_labels,
+                'protein_length': protein_length}
+
+
+@registry.register_task('secondary_structure', num_labels=3)
+class SecondaryStructureDataset(Dataset):
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False):
+
+        if split not in ('train', 'valid', 'casp12', 'ts115', 'cb513'):
+            raise ValueError(f"Unrecognized split: {split}. Must be one of "
+                             f"['train', 'valid', 'casp12', "
+                             f"'ts115', 'cb513']")
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+
+        data_path = Path(data_path)
+        data_file = f'secondary_structure/secondary_structure_{split}.lmdb'
+        self.data = dataset_factory(data_path / data_file, in_memory)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index: int):
+        item = self.data[index]
+        token_ids = self.tokenizer.encode(item['primary'])
+        input_mask = np.ones_like(token_ids)
+
+        # pad with -1s because of cls/sep tokens
+        labels = np.asarray(item['ss3'], np.int64)
+        labels = np.pad(labels, (1, 1), 'constant', constant_values=-1)
+
+        return token_ids, input_mask, labels
+
+    def collate_fn(self, batch: List[Tuple[Any, ...]]) -> Dict[str, torch.Tensor]:
+        input_ids, input_mask, ss_label = tuple(zip(*batch))
+        input_ids = torch.from_numpy(pad_sequences(input_ids, 0))
+        input_mask = torch.from_numpy(pad_sequences(input_mask, 0))
+        ss_label = torch.from_numpy(pad_sequences(ss_label, -1))
+
+        output = {'input_ids': input_ids,
+                  'input_mask': input_mask,
+                  'targets': ss_label}
+
+        return output
+
+
+@registry.register_task('trrosetta')
+class TRRosettaDataset(Dataset):
+
+    def __init__(self,
+                 data_path: Union[str, Path],
+                 split: str,
+                 tokenizer: Union[str, TAPETokenizer] = 'iupac',
+                 in_memory: bool = False,
+                 max_seqlen: int = 300):
+        if split not in ('train', 'valid'):
+            raise ValueError(
+                f"Unrecognized split: {split}. "
+                f"Must be one of ['train', 'valid']")
+        if isinstance(tokenizer, str):
+            tokenizer = TAPETokenizer(vocab=tokenizer)
+        self.tokenizer = tokenizer
+
+        data_path = Path(data_path)
+        data_path = data_path / 'trrosetta'
+        split_files = (data_path / f'{split}_files.txt').read_text().split()
+        self.data = NPZDataset(data_path / 'npz', in_memory, split_files=split_files)
+
+        self._dist_bins = np.arange(2, 20.1, 0.5)
+        self._dihedral_bins = (15 + np.arange(-180, 180, 15)) / 180 * np.pi
+        self._planar_bins = (15 + np.arange(0, 180, 15)) / 180 * np.pi
+        self._split = split
+        self.max_seqlen = max_seqlen
+        self.msa_cutoff = 0.8
+        self.penalty_coeff = 4.5
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index):
+        item = self.data[index]
+
+        msa = item['msa']
+        dist = item['dist6d']
+        omega = item['omega6d']
+        theta = item['theta6d']
+        phi = item['phi6d']
+
+        if self._split == 'train':
+            msa = self._subsample_msa(msa)
+        elif self._split == 'valid':
+            msa = msa[:20000]  # runs out of memory if msa is way too big
+        msa, dist, omega, theta, phi = self._slice_long_sequences(
+            msa, dist, omega, theta, phi)
+
+        mask = dist == 0
+
+        dist_bins = np.digitize(dist, self._dist_bins)
+        omega_bins = np.digitize(omega, self._dihedral_bins) + 1
+        theta_bins = np.digitize(theta, self._dihedral_bins) + 1
+        phi_bins = np.digitize(phi, self._planar_bins) + 1
+
+        dist_bins[mask] = 0
+        omega_bins[mask] = 0
+        theta_bins[mask] = 0
+        phi_bins[mask] = 0
+
+        dist_bins[np.diag_indices_from(dist_bins)] = -1
+
+        # input_mask = np.ones_like(msa[0])
+
+        return msa, dist_bins, omega_bins, theta_bins, phi_bins
+
+    def _slice_long_sequences(self, msa, dist, omega, theta, phi):
+        seqlen = msa.shape[1]
+        if self.max_seqlen > 0 and seqlen > self.max_seqlen:
+            start = np.random.randint(seqlen - self.max_seqlen + 1)
+            end = start + self.max_seqlen
+
+            msa = msa[:, start:end]
+            dist = dist[start:end, start:end]
+            omega = omega[start:end, start:end]
+            theta = theta[start:end, start:end]
+            phi = phi[start:end, start:end]
+
+        return msa, dist, omega, theta, phi
+
+    def _subsample_msa(self, msa):
+        num_alignments, seqlen = msa.shape
+
+        if num_alignments < 10:
+            return msa
+
+        num_sample = int(10 ** np.random.uniform(np.log10(num_alignments)) - 10)
+
+        if num_sample <= 0:
+            return msa[0][None, :]
+        elif num_sample > 20000:
+            num_sample = 20000
+
+        indices = np.random.choice(
+            msa.shape[0] - 1, size=num_sample, replace=False) + 1
+        indices = np.pad(indices, [1, 0], 'constant')  # add the sequence back in
+        return msa[indices]
+
+    def collate_fn(self, batch):
+        msa, dist_bins, omega_bins, theta_bins, phi_bins = tuple(zip(*batch))
+        # features = pad_sequences([self.featurize(msa_) for msa_ in msa], 0)
+        msa1hot = pad_sequences(
+            [F.one_hot(torch.LongTensor(msa_), 21) for msa_ in msa], 0, torch.float)
+        # input_mask = torch.FloatTensor(pad_sequences(input_mask, 0))
+        dist_bins = torch.LongTensor(pad_sequences(dist_bins, -1))
+        omega_bins = torch.LongTensor(pad_sequences(omega_bins, 0))
+        theta_bins = torch.LongTensor(pad_sequences(theta_bins, 0))
+        phi_bins = torch.LongTensor(pad_sequences(phi_bins, 0))
+
+        return {'msa1hot': msa1hot,
+                # 'input_mask': input_mask,
+                'dist': dist_bins,
+                'omega': omega_bins,
+                'theta': theta_bins,
+                'phi': phi_bins}
+
+    def featurize(self, msa):
+        msa = torch.LongTensor(msa)
+        msa1hot = F.one_hot(msa, 21).float()
+
+        seqlen = msa1hot.size(1)
+
+        weights = self.reweight(msa1hot)
+        features_1d = self.extract_features_1d(msa1hot, weights)
+        features_2d = self.extract_features_2d(msa1hot, weights)
+
+        features = torch.cat((
+            features_1d.unsqueeze(1).repeat(1, seqlen, 1),
+            features_1d.unsqueeze(0).repeat(seqlen, 1, 1),
+            features_2d), -1)
+
+        features = features.permute(2, 0, 1)
+
+        return features
+
+    def reweight(self, msa1hot):
+        # Reweight
+        seqlen = msa1hot.size(1)
+        id_min = seqlen * self.msa_cutoff
+        id_mtx = torch.tensordot(msa1hot, msa1hot, [[1, 2], [1, 2]])
+        id_mask = id_mtx > id_min
+        weights = 1.0 / id_mask.float().sum(-1)
+        return weights
+
+    def extract_features_1d(self, msa1hot, weights):
+        # 1D Features
+        seqlen = msa1hot.size(1)
+        f1d_seq = msa1hot[0, :, :20]
+
+        # msa2pssm
+        beff = weights.sum()
+        f_i = (weights[:, None, None] * msa1hot).sum(0) / beff + 1e-9
+        h_i = (-f_i * f_i.log()).sum(1, keepdims=True)
+        f1d_pssm = torch.cat((f_i, h_i), dim=1)
+
+        f1d = torch.cat((f1d_seq, f1d_pssm), dim=1)
+        f1d = f1d.view(seqlen, 42)
+        return f1d
+
+    def extract_features_2d(self, msa1hot, weights):
+        # 2D Features
+        num_alignments = msa1hot.size(0)
+        seqlen = msa1hot.size(1)
+        num_symbols = 21
+        if num_alignments == 1:
+            # No alignments, predict from sequence alone
+            f2d_dca = torch.zeros(seqlen, seqlen, 442, dtype=torch.float)
+        else:
+            # fast_dca
+
+            # covariance
+            x = msa1hot.view(num_alignments, seqlen * num_symbols)
+            num_points = weights.sum() - weights.mean().sqrt()
+            mean = (x * weights[:, None]).sum(0, keepdims=True) / num_points
+            x = (x - mean) * weights[:, None].sqrt()
+            cov = torch.matmul(x.transpose(-1, -2), x) / num_points
+
+            # inverse covariance
+            reg = torch.eye(seqlen * num_symbols) * self.penalty_coeff / weights.sum().sqrt()
+            cov_reg = cov + reg
+            inv_cov = torch.inverse(cov_reg)
+
+            x1 = inv_cov.view(seqlen, num_symbols, seqlen, num_symbols)
+            x2 = x1.permute(0, 2, 1, 3)
+            features = x2.reshape(seqlen, seqlen, num_symbols * num_symbols)
+
+            x3 = (x1[:, :-1, :, :-1] ** 2).sum((1, 3)).sqrt() * (1 - torch.eye(seqlen))
+            apc = x3.sum(0, keepdims=True) * x3.sum(1, keepdims=True) / x3.sum()
+            contacts = (x3 - apc) * (1 - torch.eye(seqlen))
+
+            f2d_dca = torch.cat([features, contacts[:, :, None]], axis=2)
+
+        return f2d_dca

tape/errors.py

@@ -0,0 +1,3 @@
+class EarlyStopping(Exception):
+    """Raised when stopping training b/c no improvement in validation loss"""
+    pass

tape/main.py

@@ -0,0 +1,269 @@
+import typing
+import os
+import logging
+import argparse
+import warnings
+import inspect
+
+
+try:
+    import apex  # noqa: F401
+    APEX_FOUND = True
+except ImportError:
+    APEX_FOUND = False
+
+from .registry import registry
+from . import training
+from . import utils
+
+CallbackList = typing.Sequence[typing.Callable]
+OutputDict = typing.Dict[str, typing.List[typing.Any]]
+
+
+logger = logging.getLogger(__name__)
+warnings.filterwarnings(  # Ignore pytorch warning about loss gathering
+    'ignore', message='Was asked to gather along dimension 0', module='torch.nn.parallel')
+
+
+def create_base_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description='Parent parser for tape functions',
+                                     add_help=False)
+    parser.add_argument('model_type', help='Base model class to run')
+    parser.add_argument('--model_config_file', default=None, type=utils.check_is_file,
+                        help='Config file for model')
+    parser.add_argument('--vocab_file', default=None,
+                        help='Pretrained tokenizer vocab file')
+    parser.add_argument('--output_dir', default='./results', type=str)
+    parser.add_argument('--no_cuda', action='store_true', help='CPU-only flag')
+    parser.add_argument('--seed', default=42, type=int, help='Random seed to use')
+    parser.add_argument('--local_rank', type=int, default=-1,
+                        help='Local rank of process in distributed training. '
+                             'Set by launch script.')
+    parser.add_argument('--tokenizer', choices=['iupac', 'unirep'],
+                        default='iupac', help='Tokenizes to use on the amino acid sequences')
+    parser.add_argument('--num_workers', default=8, type=int,
+                        help='Number of workers to use for multi-threaded data loading')
+    parser.add_argument('--log_level', default=logging.INFO,
+                        choices=['DEBUG', 'INFO', 'WARN', 'WARNING', 'ERROR',
+                                 logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR],
+                        help="log level for the experiment")
+    parser.add_argument('--debug', action='store_true', help='Run in debug mode')
+
+    return parser
+
+
+def create_train_parser(base_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description='Run Training on the TAPE datasets',
+                                     parents=[base_parser])
+    parser.add_argument('task', choices=list(registry.task_name_mapping.keys()),
+                        help='TAPE Task to train/eval on')
+    parser.add_argument('--learning_rate', default=1e-4, type=float,
+                        help='Learning rate')
+    parser.add_argument('--batch_size', default=1024, type=int,
+                        help='Batch size')
+    parser.add_argument('--data_dir', default='./data', type=utils.check_is_dir,
+                        help='Directory from which to load task data')
+    parser.add_argument('--num_train_epochs', default=10, type=int,
+                        help='Number of training epochs')
+    parser.add_argument('--num_steps_per_epoch', default=-1, type=int,
+                        help='Number of steps per epoch')
+    parser.add_argument('--num_log_iter', default=20, type=int,
+                        help='Number of training steps per log iteration')
+    parser.add_argument('--fp16', action='store_true', help='Whether to use fp16 weights')
+    parser.add_argument('--warmup_steps', default=10000, type=int,
+                        help='Number of learning rate warmup steps')
+    parser.add_argument('--gradient_accumulation_steps', default=1, type=int,
+                        help='Number of forward passes to make for each backwards pass')
+    parser.add_argument('--loss_scale', default=0, type=int,
+                        help='Loss scaling. Only used during fp16 training.')
+    parser.add_argument('--max_grad_norm', default=1.0, type=float,
+                        help='Maximum gradient norm')
+    parser.add_argument('--exp_name', default=None, type=str,
+                        help='Name to give to this experiment')
+    parser.add_argument('--from_pretrained', default=None, type=str,
+                        help='Directory containing config and pretrained model weights')
+    parser.add_argument('--log_dir', default='./logs', type=str)
+    parser.add_argument('--eval_freq', type=int, default=1,
+                        help="Frequency of eval pass. A value <= 0 means the eval pass is "
+                             "not run")
+    parser.add_argument('--save_freq', default='improvement', type=utils.int_or_str,
+                        help="How often to save the model during training. Either an integer "
+                             "frequency or the string 'improvement'")
+    parser.add_argument('--patience', default=-1, type=int,
+                        help="How many epochs without improvement to wait before ending "
+                             "training")
+    parser.add_argument('--resume_from_checkpoint', action='store_true',
+                        help="whether to resume training from the checkpoint")
+    parser.add_argument('--val_check_frac', default=1.0, type=float,
+                        help="Fraction of validation to check")                        
+    return parser
+
+
+def create_eval_parser(base_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description='Run Eval on the TAPE Datasets',
+                                     parents=[base_parser])
+    parser.add_argument('task', choices=list(registry.task_name_mapping.keys()),
+                        help='TAPE Task to train/eval on')
+    parser.add_argument('from_pretrained', type=str,
+                        help='Directory containing config and pretrained model weights')
+    parser.add_argument('--batch_size', default=1024, type=int,
+                        help='Batch size')
+    parser.add_argument('--data_dir', default='./data', type=utils.check_is_dir,
+                        help='Directory from which to load task data')
+    parser.add_argument('--metrics', default=[],
+                        help=f'Metrics to run on the result. '
+                             f'Choices: {list(registry.metric_name_mapping.keys())}',
+                        nargs='*')
+    parser.add_argument('--split', default='test', type=str,
+                        help='Which split to run on')
+    return parser
+
+
+def create_embed_parser(base_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description='Embed a set of proteins with a pretrained model',
+        parents=[base_parser])
+    parser.add_argument('data_file', type=str,
+                        help='File containing set of proteins to embed')
+    parser.add_argument('out_file', type=str,
+                        help='Name of output file')
+    parser.add_argument('from_pretrained', type=str,
+                        help='Directory containing config and pretrained model weights')
+    parser.add_argument('--batch_size', default=1024, type=int,
+                        help='Batch size')
+    parser.add_argument('--full_sequence_embed', action='store_true',
+                        help='If true, saves an embedding at every amino acid position '
+                             'in the sequence. Note that this can take a large amount '
+                             'of disk space.')
+    parser.set_defaults(task='embed')
+    return parser
+
+
+def create_distributed_parser(base_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(add_help=False, parents=[base_parser])
+    # typing.Optional arguments for the launch helper
+    parser.add_argument("--nnodes", type=int, default=1,
+                        help="The number of nodes to use for distributed "
+                             "training")
+    parser.add_argument("--node_rank", type=int, default=0,
+                        help="The rank of the node for multi-node distributed "
+                             "training")
+    parser.add_argument("--nproc_per_node", type=int, default=1,
+                        help="The number of processes to launch on each node, "
+                             "for GPU training, this is recommended to be set "
+                             "to the number of GPUs in your system so that "
+                             "each process can be bound to a single GPU.")
+    parser.add_argument("--master_addr", default="127.0.0.1", type=str,
+                        help="Master node (rank 0)'s address, should be either "
+                             "the IP address or the hostname of node 0, for "
+                             "single node multi-proc training, the "
+                             "--master_addr can simply be 127.0.0.1")
+    parser.add_argument("--master_port", default=29500, type=int,
+                        help="Master node (rank 0)'s free port that needs to "
+                             "be used for communciation during distributed "
+                             "training")
+    return parser
+
+
+def create_model_parser(base_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(add_help=False, parents=[base_parser])
+    parser.add_argument('--model_args', nargs=argparse.REMAINDER, default=None)
+    return parser
+
+def run_train(args: typing.Optional[argparse.Namespace] = None, env=None) -> None:
+    if env is not None:
+        os.environ = env
+
+    if args is None:
+        base_parser = create_base_parser()
+        train_parser = create_train_parser(base_parser)
+        model_parser = create_model_parser(train_parser)
+        args = model_parser.parse_args()
+
+    if args.gradient_accumulation_steps < 1:
+        raise ValueError(
+            f"Invalid gradient_accumulation_steps parameter: "
+            f"{args.gradient_accumulation_steps}, should be >= 1")
+
+    if (args.fp16 or args.local_rank != -1) and not APEX_FOUND:
+        raise ImportError(
+            "Please install apex from https://www.github.com/nvidia/apex "
+            "to use distributed and fp16 training.")
+
+    arg_dict = vars(args)
+    arg_names = inspect.getfullargspec(training.run_train).args
+
+    missing = set(arg_names) - set(arg_dict.keys())
+    if missing:
+        raise RuntimeError(f"Missing arguments: {missing}")
+    train_args = {name: arg_dict[name] for name in arg_names}
+
+    training.run_train(**train_args)
+
+
+def run_eval(args: typing.Optional[argparse.Namespace] = None) -> typing.Dict[str, float]:
+    if args is None:
+        base_parser = create_base_parser()
+        parser = create_eval_parser(base_parser)
+        parser = create_model_parser(parser)
+        args = parser.parse_args()
+
+    if args.from_pretrained is None:
+        raise ValueError("Must specify pretrained model")
+    if args.local_rank != -1:
+        raise ValueError("TAPE does not support distributed validation pass")
+
+    arg_dict = vars(args)
+    arg_names = inspect.getfullargspec(training.run_eval).args
+
+    missing = set(arg_names) - set(arg_dict.keys())
+    if missing:
+        raise RuntimeError(f"Missing arguments: {missing}")
+    eval_args = {name: arg_dict[name] for name in arg_names}
+
+    return training.run_eval(**eval_args)
+
+
+def run_embed(args: typing.Optional[argparse.Namespace] = None) -> None:
+    if args is None:
+        base_parser = create_base_parser()
+        parser = create_embed_parser(base_parser)
+        parser = create_model_parser(parser)
+        args = parser.parse_args()
+    if args.from_pretrained is None:
+        raise ValueError("Must specify pretrained model")
+    if args.local_rank != -1:
+        raise ValueError("TAPE does not support distributed validation pass")
+
+    arg_dict = vars(args)
+    arg_names = inspect.getfullargspec(training.run_embed).args
+
+    missing = set(arg_names) - set(arg_dict.keys())
+    if missing:
+        raise RuntimeError(f"Missing arguments: {missing}")
+    embed_args = {name: arg_dict[name] for name in arg_names}
+
+    training.run_embed(**embed_args)
+
+
+def run_train_distributed(args: typing.Optional[argparse.Namespace] = None) -> None:
+    """Runs distributed training via multiprocessing.
+    """
+    if args is None:
+        base_parser = create_base_parser()
+        distributed_parser = create_distributed_parser(base_parser)
+        distributed_train_parser = create_train_parser(distributed_parser)
+        parser = create_model_parser(distributed_train_parser)
+        args = parser.parse_args()
+
+    # Define the experiment name here, instead of dealing with barriers and communication
+    # when getting the experiment name
+    exp_name = utils.get_expname(args.exp_name, args.task, args.model_type)
+    args.exp_name = exp_name
+    utils.launch_process_group(
+        run_train, args, args.nproc_per_node, args.nnodes,
+        args.node_rank, args.master_addr, args.master_port)
+
+
+if __name__ == '__main__':
+    run_train_distributed()

tape/metrics.py

@@ -0,0 +1,48 @@
+from typing import Sequence, Union
+import numpy as np
+import scipy.stats
+
+from .registry import registry
+
+
+@registry.register_metric('mse')
+def mean_squared_error(target: Sequence[float],
+                       prediction: Sequence[float]) -> float:
+    target_array = np.asarray(target)
+    prediction_array = np.asarray(prediction)
+    return np.mean(np.square(target_array - prediction_array))
+
+
+@registry.register_metric('mae')
+def mean_absolute_error(target: Sequence[float],
+                        prediction: Sequence[float]) -> float:
+    target_array = np.asarray(target)
+    prediction_array = np.asarray(prediction)
+    return np.mean(np.abs(target_array - prediction_array))
+
+
+@registry.register_metric('spearmanr')
+def spearmanr(target: Sequence[float],
+              prediction: Sequence[float]) -> float:
+    target_array = np.asarray(target)
+    prediction_array = np.asarray(prediction)
+    return scipy.stats.spearmanr(target_array, prediction_array).correlation
+
+
+@registry.register_metric('accuracy')
+def accuracy(target: Union[Sequence[int], Sequence[Sequence[int]]],
+             prediction: Union[Sequence[float], Sequence[Sequence[float]]]) -> float:
+    if isinstance(target[0], int):
+        # non-sequence case
+        return np.mean(np.asarray(target) == np.asarray(prediction).argmax(-1))
+    else:
+        correct = 0
+        total = 0
+        for label, score in zip(target, prediction):
+            label_array = np.asarray(label)
+            pred_array = np.asarray(score).argmax(-1)
+            mask = label_array != -1
+            is_correct = label_array[mask] == pred_array[mask]
+            correct += is_correct.sum()
+            total += is_correct.size
+        return correct / total

tape/models/__init__.py

@@ -0,0 +1,28 @@
+# from .modeling_utils import ProteinConfig  # noqa: F401
+# from .modeling_utils import ProteinModel  # noqa: F401
+
+# from .modeling_bert import ProteinBertModel  # noqa: F401
+# from .modeling_bert import ProteinBertForMaskedLM  # noqa: F401
+# from .modeling_bert import ProteinBertForValuePrediction  # noqa: F401
+# from .modeling_bert import ProteinBertForSequenceClassification  # noqa: F401
+# from .modeling_bert import ProteinBertForSequenceToSequenceClassification  # noqa: F401
+# # TODO: ProteinBertForContactPrediction
+# from .modeling_resnet import ProteinResNetModel  # noqa: F401
+# from .modeling_resnet import ProteinResNetForMaskedLM  # noqa: F401
+# from .modeling_resnet import ProteinResNetForValuePrediction  # noqa: F401
+# from .modeling_resnet import ProteinResNetForSequenceClassification  # noqa: F401
+# from .modeling_resnet import ProteinResNetForSequenceToSequenceClassification  # noqa: F401
+# # TODO: ProteinResNetForContactPrediction
+# # TODO: ProteinLSTM*
+# from .modeling_unirep import UniRepModel  # noqa: F401
+# from .modeling_unirep import UniRepForLM  # noqa: F401
+# from .modeling_unirep import UniRepForValuePrediction  # noqa: F401
+# from .modeling_unirep import UniRepForSequenceClassification  # noqa: F401
+# from .modeling_unirep import UniRepForSequenceToSequenceClassification  # noqa: F401
+# # TODO: UniRepForContactPrediction
+# # TODO: Bepler*
+# from .modeling_onehot import OneHotModel  # noqa: F401
+# from .modeling_onehot import OneHotForValuePrediction  # noqa: F401
+# from .modeling_onehot import OneHotForSequenceClassification  # noqa: F401
+# from .modeling_onehot import OneHotForSequenceToSequenceClassification  # noqa: F401
+# TODO: OneHotForContactPrediction

tape/models/file_utils.py

@@ -0,0 +1,353 @@
+"""
+Utilities for working with the local dataset cache.
+This file is adapted from the huggingface transformers library at
+https://github.com/huggingface/transformers, which in turn is adapted from the AllenNLP
+library at https://github.com/allenai/allennlp
+Copyright by the AllenNLP authors.
+Note - this file goes to effort to support Python 2, but the rest of this repository does not.
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+import typing
+import sys
+import json
+import logging
+import os
+import tempfile
+import fnmatch
+from io import open
+
+import boto3
+import requests
+from botocore.exceptions import ClientError
+from tqdm import tqdm
+
+from contextlib import contextmanager
+from functools import partial, wraps
+from hashlib import sha256
+
+from filelock import FileLock
+# from tqdm.auto import tqdm
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+
+try:
+    from torch.hub import _get_torch_home
+    torch_cache_home = _get_torch_home()
+except ImportError:
+    torch_cache_home = os.path.expanduser(
+        os.getenv('TORCH_HOME', os.path.join(
+            os.getenv('XDG_CACHE_HOME', '~/.cache'), 'torch')))
+default_cache_path = os.path.join(torch_cache_home, 'protein_models')
+
+try:
+    from urllib.parse import urlparse
+except ImportError:
+    from urlparse import urlparse  # type: ignore
+
+try:
+    from pathlib import Path
+    PYTORCH_PRETRAINED_BERT_CACHE: typing.Union[str, Path] = Path(
+        os.getenv('PROTEIN_MODELS_CACHE', os.getenv(
+            'PYTORCH_PRETRAINED_BERT_CACHE', default_cache_path)))
+except (AttributeError, ImportError):
+    PYTORCH_PRETRAINED_BERT_CACHE = os.getenv('PROTEIN_MODELS_CACHE',
+                                              os.getenv('PYTORCH_PRETRAINED_BERT_CACHE',
+                                                        default_cache_path))
+
+PROTEIN_MODELS_CACHE = PYTORCH_PRETRAINED_BERT_CACHE  # Kept for backward compatibility
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+
+def get_cache():
+    return PROTEIN_MODELS_CACHE
+
+
+def get_etag(url):
+    # Get eTag to add to filename, if it exists.
+    if url.startswith("s3://"):
+        etag = s3_etag(url)
+    else:
+        try:
+            response = requests.head(url, allow_redirects=True)
+            if response.status_code != 200:
+                etag = None
+            else:
+                etag = response.headers.get("ETag")
+        except EnvironmentError:
+            etag = None
+
+    if sys.version_info[0] == 2 and etag is not None:
+        etag = etag.decode('utf-8')
+
+    return etag
+
+
+def url_to_filename(url, etag=None):
+    """
+    Convert `url` into a hashed filename in a repeatable way.
+    If `etag` is specified, append its hash to the url's, delimited
+    by a period.
+    """
+    url_bytes = url.encode('utf-8')
+    url_hash = sha256(url_bytes)
+    filename = url_hash.hexdigest()
+
+    if etag:
+        etag_bytes = etag.encode('utf-8')
+        etag_hash = sha256(etag_bytes)
+        filename += '.' + etag_hash.hexdigest()
+
+    return filename
+
+
+def filename_to_url(filename, cache_dir=None):
+    """
+    Return the url and etag (which may be ``None``) stored for `filename`.
+    Raise ``EnvironmentError`` if `filename` or its stored metadata do not exist.
+    """
+    if cache_dir is None:
+        cache_dir = PROTEIN_MODELS_CACHE
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    cache_path = os.path.join(cache_dir, filename)
+    if not os.path.exists(cache_path):
+        raise EnvironmentError("file {} not found".format(cache_path))
+
+    meta_path = cache_path + '.json'
+    if not os.path.exists(meta_path):
+        raise EnvironmentError("file {} not found".format(meta_path))
+
+    with open(meta_path, encoding="utf-8") as meta_file:
+        metadata = json.load(meta_file)
+    url = metadata['url']
+    etag = metadata['etag']
+
+    return url, etag
+
+
+def cached_path(url_or_filename, force_download=False, cache_dir=None):
+    """
+    Given something that might be a URL (or might be a local path),
+    determine which. If it's a URL, download the file and cache it, and
+    return the path to the cached file. If it's already a local path,
+    make sure the file exists and then return the path.
+
+    Args:
+        cache_dir: specify a cache directory to save the file to
+        (overwrite the default cache dir).
+        force_download: if True, re-dowload the file even if it's
+        already cached in the cache dir.
+    """
+    if cache_dir is None:
+        cache_dir = PROTEIN_MODELS_CACHE
+    if sys.version_info[0] == 3 and isinstance(url_or_filename, Path):
+        url_or_filename = str(url_or_filename)
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    parsed = urlparse(url_or_filename)
+
+    if parsed.scheme in ('http', 'https', 's3'):
+        # URL, so get it from the cache (downloading if necessary)
+        output_path = get_from_cache(url_or_filename, cache_dir, force_download)
+    elif os.path.exists(url_or_filename):
+        # File, and it exists.
+        output_path = url_or_filename
+    elif parsed.scheme == '':
+        # File, but it doesn't exist.
+        raise EnvironmentError("file {} not found".format(url_or_filename))
+    else:
+        # Something unknown
+        raise ValueError("unable to parse {} as a URL or as a local path".format(
+            url_or_filename))
+
+    return output_path
+
+
+def split_s3_path(url):
+    """Split a full s3 path into the bucket name and path."""
+    parsed = urlparse(url)
+    if not parsed.netloc or not parsed.path:
+        raise ValueError("bad s3 path {}".format(url))
+    bucket_name = parsed.netloc
+    s3_path = parsed.path
+    # Remove '/' at beginning of path.
+    if s3_path.startswith("/"):
+        s3_path = s3_path[1:]
+    return bucket_name, s3_path
+
+
+def s3_request(func):
+    """
+    Wrapper function for s3 requests in order to create more helpful error
+    messages.
+    """
+
+    @wraps(func)
+    def wrapper(url, *args, **kwargs):
+        try:
+            return func(url, *args, **kwargs)
+        except ClientError as exc:
+            if int(exc.response["Error"]["Code"]) == 404:
+                raise EnvironmentError("file {} not found".format(url))
+            else:
+                raise
+
+    return wrapper
+
+
+@s3_request
+def s3_etag(url):
+    """Check ETag on S3 object."""
+    s3_resource = boto3.resource("s3")
+    bucket_name, s3_path = split_s3_path(url)
+    s3_object = s3_resource.Object(bucket_name, s3_path)
+    return s3_object.e_tag
+
+
+@s3_request
+def s3_get(url, temp_file):
+    """Pull a file directly from S3."""
+    s3_resource = boto3.resource("s3")
+    bucket_name, s3_path = split_s3_path(url)
+    s3_resource.Bucket(bucket_name).download_fileobj(s3_path, temp_file)
+
+
+def http_get(url, temp_file):
+    req = requests.get(url, stream=True)
+    content_length = req.headers.get('Content-Length')
+    total = int(content_length) if content_length is not None else None
+    progress = tqdm(unit="B", total=total)
+    for chunk in req.iter_content(chunk_size=1024):
+        if chunk:  # filter out keep-alive new chunks
+            progress.update(len(chunk))
+            temp_file.write(chunk)
+    progress.close()
+
+
+def get_from_cache(url, cache_dir=None, force_download=False, resume_download=False):
+    """
+    Given a URL, look for the corresponding dataset in the local cache.
+    If it's not there, download it. Then return the path to the cached file.
+    """
+    if cache_dir is None:
+        cache_dir = PROTEIN_MODELS_CACHE
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+    if sys.version_info[0] == 2 and not isinstance(cache_dir, str):
+        cache_dir = str(cache_dir)
+
+    if not os.path.exists(cache_dir):
+        os.makedirs(cache_dir)
+
+    # Get eTag to add to filename, if it exists.
+    if url.startswith("s3://"):
+        etag = s3_etag(url)
+    else:
+        try:
+            response = requests.head(url, allow_redirects=True)
+            if response.status_code != 200:
+                etag = None
+            else:
+                etag = response.headers.get("ETag")
+        except EnvironmentError:
+            etag = None
+
+    if sys.version_info[0] == 2 and etag is not None:
+        etag = etag.decode('utf-8')
+    filename = url_to_filename(url, etag)
+
+    # get cache path to put the file
+    cache_path = os.path.join(cache_dir, filename)
+
+    if os.path.exists(cache_path) and etag is None:
+        return cache_path
+
+    # If we don't have a connection (etag is None) and can't identify the file
+    # try to get the last downloaded one
+    if not os.path.exists(cache_path) and etag is None:
+        matching_files = fnmatch.filter(os.listdir(cache_dir), filename + '.*')
+        matching_files = list(filter(lambda s: not s.endswith('.json'), matching_files))
+        if matching_files:
+            cache_path = os.path.join(cache_dir, matching_files[-1])
+
+    # From now on, etag is not None
+    if os.path.exists(cache_path) and not force_download:
+        return cache_path
+
+    # Prevent parallel downloads of the same file with a lock.
+    lock_path = cache_path + ".lock"
+    with FileLock(lock_path):
+
+        # If the download just completed while the lock was activated.
+        if os.path.exists(cache_path) and not force_download:
+            # Even if returning early like here, the lock will be released.
+            return cache_path
+
+        if resume_download:
+            incomplete_path = cache_path + ".incomplete"
+
+            @contextmanager
+            def _resumable_file_manager():
+                with open(incomplete_path, "a+b") as f:
+                    yield f
+
+            temp_file_manager = _resumable_file_manager
+        else:
+            temp_file_manager = partial(tempfile.NamedTemporaryFile, dir=cache_dir,
+                                        delete=False)
+        # Download to temporary file, then copy to cache dir once finished.
+        # Otherwise you get corrupt cache entries if the download gets interrupted.
+        with temp_file_manager() as temp_file:
+            logger.info("%s not in cache or force_download=True, download to %s",
+                        url, temp_file.name)
+
+            http_get(url, temp_file)
+
+        logger.info("storing %s in cache at %s", url, cache_path)
+        os.replace(temp_file.name, cache_path)
+
+        logger.info("creating metadata file for %s", cache_path)
+        meta = {"url": url, "etag": etag}
+        meta_path = cache_path + ".json"
+        with open(meta_path, "w") as meta_file:
+            json.dump(meta, meta_file)
+    '''
+    if not os.path.exists(cache_path):
+        # Download to temporary file, then copy to cache dir once finished.
+        # Otherwise you get corrupt cache entries if the download gets interrupted.
+        with tempfile.NamedTemporaryFile() as temp_file:
+            logger.info("%s not found in cache, downloading to %s", url, temp_file.name)
+
+            # GET file object
+            if url.startswith("s3://"):
+                s3_get(url, temp_file)
+            else:
+                http_get(url, temp_file)
+
+            # we are copying the file before closing it, so flush to avoid truncation
+            temp_file.flush()
+            # shutil.copyfileobj() starts at the current position, so go to the start
+            temp_file.seek(0)
+
+            logger.info("copying %s to cache at %s", temp_file.name, cache_path)
+            with open(cache_path, 'wb') as cache_file:
+                shutil.copyfileobj(temp_file, cache_file)
+
+            logger.info("creating metadata file for %s", cache_path)
+            meta = {'url': url, 'etag': etag}
+            meta_path = cache_path + '.json'
+            with open(meta_path, 'w') as meta_file:
+                output_string = json.dumps(meta)
+                if sys.version_info[0] == 2 and isinstance(output_string, str):
+                    # The beauty of python 2
+                    output_string = unicode(output_string, 'utf-8')  # noqa: F821
+                meta_file.write(output_string)
+
+            logger.info("removing temp file %s", temp_file.name)
+    '''
+    return cache_path

tape/models/modeling_autoencoder.py

@@ -0,0 +1,316 @@
+import typing
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .modeling_utils import ProteinConfig
+from .modeling_utils import ProteinModel
+from .modeling_utils import get_activation_fn
+from .modeling_utils import MLMHead
+from .modeling_utils import LayerNorm
+from .modeling_utils import ValuePredictionHead
+from .modeling_utils import SequenceClassificationHead
+from .modeling_utils import SequenceToSequenceClassificationHead
+from .modeling_utils import PairwiseContactPredictionHead
+from ..registry import registry
+
+logger = logging.getLogger(__name__)
+
+RESNET_PRETRAINED_CONFIG_ARCHIVE_MAP: typing.Dict[str, str] = {}
+RESNET_PRETRAINED_MODEL_ARCHIVE_MAP: typing.Dict[str, str] = {}
+
+
+class ProteinAEConfig(ProteinConfig):
+    pretrained_config_archive_map = RESNET_PRETRAINED_CONFIG_ARCHIVE_MAP
+
+    def __init__(self,
+                 vocab_size: int = 30,
+                 hidden_size: int = 512,
+                 num_hidden_layers: int = 30,
+                 hidden_act: str = "gelu",
+                 hidden_dropout_prob: float = 0.1,
+                 initializer_range: float = 0.02,
+                 layer_norm_eps: float = 1e-12,
+                 temporal_pooling: str = 'attention',
+                 freeze_embedding: bool = False,
+                 max_size: int = 3000,
+                 latent_size: int = 1024,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.num_hidden_layers = num_hidden_layers
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.temporal_pooling = temporal_pooling
+        self.freeze_embedding = freeze_embedding
+        self.max_size = max_size
+        self.latent_size = latent_size
+
+
+class MaskedConv1d(nn.Conv1d):
+
+    def forward(self, x, input_mask=None):
+        if input_mask is not None:
+            x = x * input_mask
+        return super().forward(x)
+
+
+class ProteinResNetLayerNorm(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.norm = LayerNorm(config.hidden_size)
+
+    def forward(self, x):
+        return self.norm(x.transpose(1, 2)).transpose(1, 2)
+
+
+class ProteinResNetBlock(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.conv1 = MaskedConv1d(
+            config.hidden_size, config.hidden_size, 3, padding=1, bias=False)
+        # self.bn1 = nn.BatchNorm1d(config.hidden_size)
+        self.bn1 = ProteinResNetLayerNorm(config)
+        self.conv2 = MaskedConv1d(
+            config.hidden_size, config.hidden_size, 3, padding=1, bias=False)
+        # self.bn2 = nn.BatchNorm1d(config.hidden_size)
+        self.bn2 = ProteinResNetLayerNorm(config)
+        self.activation_fn = get_activation_fn(config.hidden_act)
+
+    def forward(self, x, input_mask=None):
+        identity = x
+
+        out = self.conv1(x, input_mask)
+        out = self.bn1(out)
+        out = self.activation_fn(out)
+
+        out = self.conv2(out, input_mask)
+        out = self.bn2(out)
+
+        out += identity
+        out = self.activation_fn(out)
+
+        return out
+
+
+class ProteinResNetEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings.
+    """
+    def __init__(self, config):
+        super().__init__()
+        embed_dim = config.hidden_size
+        self.word_embeddings = nn.Embedding(config.vocab_size, embed_dim, padding_idx=0)
+        inverse_frequency = 1 / (10000 ** (torch.arange(0.0, embed_dim, 2.0) / embed_dim))
+        self.register_buffer('inverse_frequency', inverse_frequency)
+
+        self.layer_norm = LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, input_ids):
+        words_embeddings = self.word_embeddings(input_ids)
+
+        seq_length = input_ids.size(1)
+        position_ids = torch.arange(
+            seq_length - 1, -1, -1.0,
+            dtype=words_embeddings.dtype,
+            device=words_embeddings.device)
+        sinusoidal_input = torch.ger(position_ids, self.inverse_frequency)
+        position_embeddings = torch.cat([sinusoidal_input.sin(), sinusoidal_input.cos()], -1)
+        position_embeddings = position_embeddings.unsqueeze(0)
+
+        embeddings = words_embeddings + position_embeddings
+        embeddings = self.layer_norm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class ResNetEncoder(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.output_hidden_states = config.output_hidden_states
+        self.encoder = nn.ModuleList(
+            [ProteinResNetBlock(config) for _ in range(config.num_hidden_layers)])
+
+        self.decoder = nn.ModuleList(
+            [ProteinResNetBlock(config) for _ in range(config.num_hidden_layers)])
+
+        self.bottleneck1 = nn.Linear(93*config.hidden_size, config.latent_size)
+        self.bottleneck2 = nn.Linear(config.latent_size, 94*config.hidden_size)
+
+    def forward(self, hidden_states, input_mask=None):
+        for i, layer_module in enumerate(self.encoder):
+            hidden_states = layer_module(hidden_states)
+            if i != 0 and i % 5 == 0:
+                hidden_states = nn.functional.avg_pool1d(hidden_states, 2, stride=2)
+
+        bs = hidden_states.shape[0]
+        latents = self.bottleneck1(hidden_states.reshape(bs, -1))
+        hidden_states = self.bottleneck2(latents).reshape(bs, -1, 94)
+
+
+        for i, layer_module in enumerate(self.decoder):
+            if i != 0 and i % 5 == 0:
+                hidden_states = nn.functional.interpolate(hidden_states, scale_factor=2)
+            hidden_states = layer_module(hidden_states)
+
+        hidden_states = hidden_states[:,:,:self.config.max_size]
+        outputs = (hidden_states, latents)
+
+        return outputs
+
+
+class ProteinAEAbstractModel(ProteinModel):
+    """ An abstract class to handle weights initialization and
+        a simple interface for dowloading and loading pretrained models.
+    """
+    config_class = ProteinAEConfig
+    base_model_prefix = "ae"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Conv1d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+
+@registry.register_task_model('embed', 'autoencoder')
+class ProteinResNetModel(ProteinAEAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = ProteinResNetEmbeddings(config)
+        self.encoder = ResNetEncoder(config)
+
+        self.init_weights()
+
+    def forward(self,
+                input_ids,
+                input_mask=None):
+        pre_pad_shape = input_ids.shape[1]
+        if pre_pad_shape >= self.config.max_size:
+            input_ids = input_ids[:,:self.config.max_size]
+            if not input_mask is None:
+                input_mask = input_mask[:,:self.config.max_size]
+        else:
+            input_ids = F.pad(input_ids, (0, self.config.max_size - pre_pad_shape))
+            if not input_mask is None:
+                input_mask = F.pad(input_mask, (0, self.config.max_size - pre_pad_shape))
+        assert input_ids.shape[1] == self.config.max_size
+        
+        if input_mask is not None and torch.any(input_mask != 1):
+            extended_input_mask = input_mask.unsqueeze(2)
+            # fp16 compatibility
+            extended_input_mask = extended_input_mask.to(
+                dtype=next(self.parameters()).dtype)
+        else:
+            extended_input_mask = None
+
+        embedding_output = self.embeddings(input_ids)
+        embedding_output = embedding_output.transpose(1, 2)
+        if extended_input_mask is not None:
+            extended_input_mask = extended_input_mask.transpose(1, 2)
+        sequence_output, pooled_output = self.encoder(embedding_output, extended_input_mask)
+        sequence_output = sequence_output.transpose(1, 2).contiguous()
+        return sequence_output, pooled_output
+
+@registry.register_task_model('beta_lactamase', 'autoencoder')
+@registry.register_task_model('language_modeling', 'autoencoder')
+class ProteinResNetForMaskedLM(ProteinAEAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.resnet = ProteinResNetModel(config)
+        self.mlm = MLMHead(
+            config.hidden_size, config.vocab_size, config.hidden_act, config.layer_norm_eps,
+            ignore_index=-1)
+ 
+        self.init_weights()
+        self.tie_weights()
+
+    def tie_weights(self):
+        """ Make sure we are sharing the input and output embeddings.
+            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
+        """
+        self._tie_or_clone_weights(self.mlm.decoder,
+                                   self.resnet.embeddings.word_embeddings)
+
+    def forward(self,
+                input_ids,
+                input_mask=None,
+                targets=None):
+        pre_pad_shape = input_ids.shape[1]
+        if targets is not None:
+            targets = targets[:,:self.config.max_size]
+        
+        outputs = self.resnet(input_ids, input_mask=input_mask)      
+        outputs = self.mlm(outputs[0][:,:pre_pad_shape,:], targets) + (outputs[1],)
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('fluorescence', 'autoencoder')
+@registry.register_task_model('stability', 'autoencoder')
+class ProteinResNetForValuePrediction(ProteinAEAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.resnet = ProteinResNetModel(config)
+        self.predict = ValuePredictionHead(config.hidden_size)
+        self.freeze_embedding = config.freeze_embedding
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+        if self.freeze_embedding:
+            self.resnet.train(False)
+
+        outputs = self.resnet(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('remote_homology', 'autoencoder')
+class ProteinResNetForSequenceClassification(ProteinAEAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.resnet = ProteinResNetModel(config)
+        self.classify = SequenceClassificationHead(config.hidden_size, config.num_labels)
+        self.freeze_embedding = config.freeze_embedding
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+        if self.freeze_embedding:
+            self.resnet.train(False)
+
+        outputs = self.resnet(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs

tape/models/modeling_bert.py

@@ -0,0 +1,606 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+# Modified by Roshan Rao
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch BERT model. """
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+import math
+
+import torch
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from .modeling_utils import ProteinConfig
+from .modeling_utils import ProteinModel
+from .modeling_utils import prune_linear_layer
+from .modeling_utils import get_activation_fn
+from .modeling_utils import LayerNorm
+from .modeling_utils import MLMHead
+from .modeling_utils import ValuePredictionHead
+from .modeling_utils import SequenceClassificationHead
+from .modeling_utils import SequenceToSequenceClassificationHead
+from .modeling_utils import PairwiseContactPredictionHead
+from ..registry import registry
+
+logger = logging.getLogger(__name__)
+
+URL_PREFIX = "https://s3.amazonaws.com/proteindata/pytorch-models/"
+BERT_PRETRAINED_MODEL_ARCHIVE_MAP = {
+    'bert-base': URL_PREFIX + "bert-base-pytorch_model.bin",
+}
+BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    'bert-base': URL_PREFIX + "bert-base-config.json"
+}
+
+
+class ProteinBertConfig(ProteinConfig):
+    r"""
+        :class:`~pytorch_transformers.ProteinBertConfig` is the configuration class to store the
+        configuration of a `ProteinBertModel`.
+
+
+        Arguments:
+            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in
+                `ProteinBertModel`.
+            hidden_size: Size of the encoder layers and the pooler layer.
+            num_hidden_layers: Number of hidden layers in the ProteinBert encoder.
+            num_attention_heads: Number of attention heads for each attention layer in
+                the ProteinBert encoder.
+            intermediate_size: The size of the "intermediate" (i.e., feed-forward)
+                layer in the ProteinBert encoder.
+            hidden_act: The non-linear activation function (function or string) in the
+                encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
+            hidden_dropout_prob: The dropout probabilitiy for all fully connected
+                layers in the embeddings, encoder, and pooler.
+            attention_probs_dropout_prob: The dropout ratio for the attention
+                probabilities.
+            max_position_embeddings: The maximum sequence length that this model might
+                ever be used with. Typically set this to something large just in case
+                (e.g., 512 or 1024 or 2048).
+            type_vocab_size: The vocabulary size of the `token_type_ids` passed into
+                `ProteinBertModel`.
+            initializer_range: The sttdev of the truncated_normal_initializer for
+                initializing all weight matrices.
+            layer_norm_eps: The epsilon used by LayerNorm.
+    """
+    pretrained_config_archive_map = BERT_PRETRAINED_CONFIG_ARCHIVE_MAP
+
+    def __init__(self,
+                 vocab_size: int = 30,
+                 hidden_size: int = 768,
+                 num_hidden_layers: int = 12,
+                 num_attention_heads: int = 12,
+                 intermediate_size: int = 3072,
+                 hidden_act: str = "gelu",
+                 hidden_dropout_prob: float = 0.1,
+                 attention_probs_dropout_prob: float = 0.1,
+                 max_position_embeddings: int = 8096,
+                 type_vocab_size: int = 2,
+                 initializer_range: float = 0.02,
+                 layer_norm_eps: float = 1e-12,
+                 temporal_pooling: str = 'attention',
+                 freeze_embedding: bool = False,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.temporal_pooling = temporal_pooling
+        self.freeze_embedding = freeze_embedding
+
+
+class ProteinBertEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings.
+    """
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, padding_idx=0)
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be
+        # able to load any TensorFlow checkpoint file
+        self.LayerNorm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, input_ids, token_type_ids=None, position_ids=None):
+        seq_length = input_ids.size(1)
+        if position_ids is None:
+            position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
+            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        words_embeddings = self.word_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = words_embeddings + position_embeddings + token_type_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class ProteinBertSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (config.hidden_size, config.num_attention_heads))
+        self.output_attentions = config.output_attentions
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(self, hidden_states, attention_mask):
+        mixed_query_layer = self.query(hidden_states)
+        mixed_key_layer = self.key(hidden_states)
+        mixed_value_layer = self.value(hidden_states)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+        key_layer = self.transpose_for_scores(mixed_key_layer)
+        value_layer = self.transpose_for_scores(mixed_value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        # Apply the attention mask is (precomputed for all layers in
+        # ProteinBertModel forward() function)
+        attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original ProteinBert paper.
+        attention_probs = self.dropout(attention_probs)
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) \
+            if self.output_attentions else (context_layer,)
+        return outputs
+
+
+class ProteinBertSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class ProteinBertAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self = ProteinBertSelfAttention(config)
+        self.output = ProteinBertSelfOutput(config)
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        mask = torch.ones(self.self.num_attention_heads, self.self.attention_head_size)
+        for head in heads:
+            mask[head] = 0
+        mask = mask.view(-1).contiguous().eq(1)
+        index = torch.arange(len(mask))[mask].long()
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+        # Update hyper params
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+
+    def forward(self, input_tensor, attention_mask):
+        self_outputs = self.self(input_tensor, attention_mask)
+        attention_output = self.output(self_outputs[0], input_tensor)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+class ProteinBertIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = get_activation_fn(config.hidden_act)
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class ProteinBertOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class ProteinBertLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention = ProteinBertAttention(config)
+        self.intermediate = ProteinBertIntermediate(config)
+        self.output = ProteinBertOutput(config)
+
+    def forward(self, hidden_states, attention_mask):
+        attention_outputs = self.attention(hidden_states, attention_mask)
+        attention_output = attention_outputs[0]
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        outputs = (layer_output,) + attention_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+class ProteinBertEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.output_attentions = config.output_attentions
+        self.output_hidden_states = config.output_hidden_states
+        self.layer = nn.ModuleList(
+            [ProteinBertLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def run_function(self, start, chunk_size):
+        def custom_forward(hidden_states, attention_mask):
+            all_hidden_states = ()
+            all_attentions = ()
+            chunk_slice = slice(start, start + chunk_size)
+            for layer in self.layer[chunk_slice]:
+                if self.output_hidden_states:
+                    all_hidden_states = all_hidden_states + (hidden_states,)
+                layer_outputs = layer(hidden_states, attention_mask)
+                hidden_states = layer_outputs[0]
+
+                if self.output_attentions:
+                    all_attentions = all_attentions + (layer_outputs[1],)
+
+            if self.output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            outputs = (hidden_states,)
+            if self.output_hidden_states:
+                outputs = outputs + (all_hidden_states,)
+            if self.output_attentions:
+                outputs = outputs + (all_attentions,)
+            return outputs
+
+        return custom_forward
+
+    def forward(self, hidden_states, attention_mask, chunks=None):
+        all_hidden_states = ()
+        all_attentions = ()
+
+        if chunks is not None:
+            assert isinstance(chunks, int)
+            chunk_size = (len(self.layer) + chunks - 1) // chunks
+            for start in range(0, len(self.layer), chunk_size):
+                outputs = checkpoint(self.run_function(start, chunk_size),
+                                     hidden_states, attention_mask)
+                if self.output_hidden_states:
+                    all_hidden_states = all_hidden_states + outputs[1]
+                if self.output_attentions:
+                    all_attentions = all_attentions + outputs[-1]
+                hidden_states = outputs[0]
+        else:
+            for i, layer_module in enumerate(self.layer):
+                if self.output_hidden_states:
+                    all_hidden_states = all_hidden_states + (hidden_states,)
+
+                layer_outputs = layer_module(hidden_states, attention_mask)
+                hidden_states = layer_outputs[0]
+
+                if self.output_attentions:
+                    all_attentions = all_attentions + (layer_outputs[1],)
+
+            # Add last layer
+            if self.output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            outputs = (hidden_states,)
+            if self.output_hidden_states:
+                outputs = outputs + (all_hidden_states,)
+            if self.output_attentions:
+                outputs = outputs + (all_attentions,)
+        return outputs  # outputs, (hidden states), (attentions)
+
+
+class ProteinBertPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+        self.temporal_pooling = config.temporal_pooling
+        self._la_w1 = nn.Conv1d(config.hidden_size, int(config.hidden_size/2), 5, padding=2)
+        self._la_w2 = nn.Conv1d(config.hidden_size, int(config.hidden_size/2), 5, padding=2)
+        self._la_mlp = nn.Linear(config.hidden_size, config.hidden_size)
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        if self.temporal_pooling == 'mean':
+            return hidden_states.mean(dim=1)
+        if self.temporal_pooling == 'max':
+            return hidden_states.max(dim=1)
+        if self.temporal_pooling == 'concat':
+            _temp = hidden_states.reshape(hidden_states.shape[0], -1)
+            return torch.nn.functional.pad(_temp, (0, 2048 - _temp.shape[1]))
+        if self.temporal_pooling == 'topmax':
+            val, _ = torch.topk(hidden_states, k=5, dim=1)
+            return val.mean(dim=1)
+        if self.temporal_pooling == 'light_attention':
+            _temp = hidden_states.permute(0,2,1)
+            a = self._la_w1(_temp).softmax(dim=-1)
+            v = self._la_w2(_temp)
+            v_max = v.max(dim=-1).values
+            v_sum = (a * v).sum(dim=-1)
+            return self._la_mlp(torch.cat([v_max, v_sum], dim=1))
+
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class ProteinBertAbstractModel(ProteinModel):
+    """ An abstract class to handle weights initialization and
+        a simple interface for dowloading and loading pretrained models.
+    """
+    config_class = ProteinBertConfig
+    pretrained_model_archive_map = BERT_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+@registry.register_task_model('embed', 'transformer')
+class ProteinBertModel(ProteinBertAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = ProteinBertEmbeddings(config)
+        self.encoder = ProteinBertEncoder(config)
+        self.pooler = ProteinBertPooler(config)
+
+        self.init_weights()
+
+    def _resize_token_embeddings(self, new_num_tokens):
+        old_embeddings = self.embeddings.word_embeddings
+        new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens)
+        self.embeddings.word_embeddings = new_embeddings
+        return self.embeddings.word_embeddings
+
+    def _prune_heads(self, heads_to_prune):
+        """ Prunes heads of the model.
+            heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
+            See base class ProteinModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def forward(self,
+                input_ids,
+                input_mask=None):
+        if input_mask is None:
+            input_mask = torch.ones_like(input_ids)
+
+        # We create a 3D attention mask from a 2D tensor mask.
+        # Sizes are [batch_size, 1, 1, to_seq_length]
+        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
+        # this attention mask is more simple than the triangular masking of causal attention
+        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
+        extended_attention_mask = input_mask.unsqueeze(1).unsqueeze(2)
+
+        # Since input_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(
+            dtype=torch.float32)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+
+        embedding_output = self.embeddings(input_ids)
+        encoder_outputs = self.encoder(embedding_output,
+                                       extended_attention_mask,
+                                       chunks=None)
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output)
+
+        # add hidden_states and attentions if they are here
+        outputs = (sequence_output, pooled_output,) + encoder_outputs[1:]
+        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)
+
+
+@registry.register_task_model('masked_language_modeling', 'transformer')
+class ProteinBertForMaskedLM(ProteinBertAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = ProteinBertModel(config)
+        self.mlm = MLMHead(
+            config.hidden_size, config.vocab_size, config.hidden_act, config.layer_norm_eps,
+            ignore_index=-1)
+
+        self.init_weights()
+        self.tie_weights()
+
+    def tie_weights(self):
+        """ Make sure we are sharing the input and output embeddings.
+            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
+        """
+        self._tie_or_clone_weights(self.mlm.decoder,
+                                   self.bert.embeddings.word_embeddings)
+
+    def forward(self,
+                input_ids,
+                input_mask=None,
+                targets=None):
+
+        outputs = self.bert(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        # add hidden states and attention if they are here
+        outputs = self.mlm(sequence_output, targets) + outputs[:2]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('fluorescence', 'transformer')
+@registry.register_task_model('stability', 'transformer')
+class ProteinBertForValuePrediction(ProteinBertAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = ProteinBertModel(config)
+        self.predict = ValuePredictionHead(config.hidden_size)
+        self.freeze_embedding = config.freeze_embedding
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+        if self.freeze_embedding:
+            self.bert.train(False)
+        outputs = self.bert(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('remote_homology', 'transformer')
+class ProteinBertForSequenceClassification(ProteinBertAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = ProteinBertModel(config)
+        self.classify = SequenceClassificationHead(
+            config.hidden_size, config.num_labels)
+        self.freeze_embedding = config.freeze_embedding
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+        if self.freeze_embedding:
+            self.bert.train(False)
+        outputs = self.bert(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+
+        outputs = self.classify(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('secondary_structure', 'transformer')
+class ProteinBertForSequenceToSequenceClassification(ProteinBertAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = ProteinBertModel(config)
+        self.classify = SequenceToSequenceClassificationHead(
+            config.hidden_size, config.num_labels, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.bert(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(sequence_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('contact_prediction', 'transformer')
+class ProteinBertForContactPrediction(ProteinBertAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = ProteinBertModel(config)
+        self.predict = PairwiseContactPredictionHead(config.hidden_size, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, protein_length, input_mask=None, targets=None):
+
+        outputs = self.bert(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(sequence_output, protein_length, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs

tape/models/modeling_bottleneck.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tape import ProteinModel, ProteinConfig
+from tape.models.modeling_utils import SequenceToSequenceClassificationHead
+from tape.registry import registry
+from .modeling_utils import LayerNorm, MLMHead
+from .modeling_bert import ProteinBertModel, ProteinBertConfig
+from .modeling_lstm import ProteinLSTMModel, ProteinLSTMConfig
+from .modeling_resnet import ProteinResNetModel, ProteinResNetConfig
+
+
+class BottleneckConfig(ProteinConfig):
+    def __init__(self,
+                 hidden_size: int = 1024,
+                 max_size: int = 300,
+                 backend_name: str = 'resnet',
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.max_size = max_size
+        self.backend_name = backend_name
+
+
+class BottleneckAbstractModel(ProteinModel):
+    """ All your models will inherit from this one - it's used to define the
+        config_class of the model set and also to define the base_model_prefix.
+        This is used to allow easy loading/saving into different models.
+    """
+    config_class = BottleneckConfig
+    base_model_prefix = 'bottleneck'
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.Conv1d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
+            if module.bias is not None:
+                module.bias.data.zero_()
+        # elif isinstance(module, ProteinResNetBlock):
+            # nn.init.constant_(module.bn2.weight, 0)
+
+@registry.register_task_model('embed', 'bottleneck')
+class ProteinBottleneckModel(BottleneckAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        if config.backend_name == 'resnet':
+            config = ProteinResNetConfig()
+            self.backbone1 = ProteinResNetModel(config)
+        elif config.backend_name == 'transformer':
+            config = ProteinBertConfig()
+            self.backbone1 = ProteinBertModel(config)
+        elif config.backend_name == 'lstm':
+            config = ProteinLSTMConfig(hidden_size=256)
+            self.backbone1 = ProteinLSTMModel(config)
+            config.hidden_size = config.hidden_size * 2
+        else:
+            raise ValueError('Somethings wrong')
+        self.linear1 = nn.Linear(self.config.max_size*config.hidden_size, self.config.hidden_size)
+        self.linear2 = nn.Linear(self.config.hidden_size, self.config.max_size*config.hidden_size)
+
+    def forward(self, input_ids, input_mask=None):
+        pre_pad_shape = input_ids.shape[1]
+        if pre_pad_shape >= self.config.max_size:
+            input_ids = input_ids[:,:self.config.max_size]
+            if not input_mask is None:
+                input_mask = input_mask[:,:self.config.max_size]
+        else:
+            input_ids = F.pad(input_ids, (0, self.config.max_size - pre_pad_shape))
+            if not input_mask is None:
+                input_mask = F.pad(input_mask, (0, self.config.max_size - pre_pad_shape))
+        assert input_ids.shape[1] == self.config.max_size
+
+        output = self.backbone1(input_ids, input_mask)
+        sequence_output = output[0]
+        pre_shape = sequence_output.shape
+        embeddings = self.linear1(sequence_output.reshape(sequence_output.shape[0], -1))
+        sequence_output = self.linear2(embeddings).reshape(*pre_shape)
+        sequence_output = sequence_output[:,:pre_pad_shape]
+        outputs = (sequence_output, embeddings) + output[2:]
+        return outputs
+
+@registry.register_task_model('beta_lactamase', 'bottleneck')
+@registry.register_task_model('masked_language_modeling', 'bottleneck')
+@registry.register_task_model('language_modeling', 'bottleneck')
+class ProteinBottleneckForPretraining(BottleneckAbstractModel):
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone1 = ProteinBottleneckModel(config)
+        
+        if config.backend_name == 'resnet':
+            config = ProteinResNetConfig()
+            self.backbone2 = MLMHead(config.hidden_size, config.vocab_size, config.hidden_act, 
+                                     config.layer_norm_eps, ignore_index=-1)
+        elif config.backend_name == 'transformer':
+            config = ProteinBertConfig()
+            self.backbone2 = MLMHead(config.hidden_size, config.vocab_size, config.hidden_act, 
+                                     config.layer_norm_eps, ignore_index=-1)
+        elif config.backend_name == 'lstm':
+            config = ProteinLSTMConfig(hidden_size=256)
+            self.backbone2 = nn.Linear(config.hidden_size, config.vocab_size)
+            config.hidden_size = config.hidden_size * 2
+        else:
+            raise ValueError('Somethings wrong')
+
+    def forward(self,
+                input_ids,
+                input_mask=None,
+                targets=None):
+        if input_ids.shape[1]>self.config.max_size:
+            targets = targets[:,:self.config.max_size]
+            
+        outputs = self.backbone1(input_ids, input_mask)
+        sequence_output = outputs[0]
+        if self.config.backend_name == 'resnet' or self.config.backend_name == 'transformer':
+            outputs = self.backbone2(sequence_output, targets) + outputs[2:]
+        elif self.config.backend_name == 'lstm':
+            sequence_output, pooled_output = outputs[:2]
+
+            forward_prediction, reverse_prediction = sequence_output.chunk(2, -1)
+            forward_prediction = F.pad(forward_prediction[:, :-1], [0, 0, 1, 0])
+            reverse_prediction = F.pad(reverse_prediction[:, 1:], [0, 0, 0, 1])
+            prediction_scores = \
+                self.backbone2(forward_prediction) + self.backbone2(reverse_prediction)
+            prediction_scores = prediction_scores.contiguous()
+
+            # add hidden states and if they are here
+            outputs = (prediction_scores,) + outputs[2:]
+
+            if targets is not None:
+                loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
+                lm_loss = loss_fct(
+                    prediction_scores.view(-1, 30), targets.view(-1))
+                outputs = (lm_loss,) + outputs           
+
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs

tape/models/modeling_lstm.py

@@ -0,0 +1,335 @@
+import logging
+import typing
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .modeling_utils import ProteinConfig
+from .modeling_utils import ProteinModel
+from .modeling_utils import ValuePredictionHead
+from .modeling_utils import SequenceClassificationHead
+from .modeling_utils import SequenceToSequenceClassificationHead
+from .modeling_utils import PairwiseContactPredictionHead
+from ..registry import registry
+
+logger = logging.getLogger(__name__)
+
+
+URL_PREFIX = "https://s3.amazonaws.com/proteindata/pytorch-models/"
+LSTM_PRETRAINED_CONFIG_ARCHIVE_MAP: typing.Dict[str, str] = {}
+LSTM_PRETRAINED_MODEL_ARCHIVE_MAP: typing.Dict[str, str] = {}
+
+
+class ProteinLSTMConfig(ProteinConfig):
+    pretrained_config_archive_map = LSTM_PRETRAINED_CONFIG_ARCHIVE_MAP
+
+    def __init__(self,
+                 vocab_size: int = 30,
+                 input_size: int = 128,
+                 hidden_size: int = 1024,
+                 num_hidden_layers: int = 3,
+                 hidden_dropout_prob: float = 0.1,
+                 initializer_range: float = 0.02,
+                 temporal_pooling: str = 'attention',
+                 freeze_embedding: bool = False,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.initializer_range = initializer_range
+        self.temporal_pooling = temporal_pooling
+        self.freeze_embedding = freeze_embedding
+
+
+class ProteinLSTMLayer(nn.Module):
+
+    def __init__(self, input_size: int, hidden_size: int, dropout: float = 0.):
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
+
+    def forward(self, inputs):
+        inputs = self.dropout(inputs)
+        self.lstm.flatten_parameters()
+        return self.lstm(inputs)
+
+
+class ProteinLSTMPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.scalar_reweighting = nn.Linear(2 * config.num_hidden_layers, 1)
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+        self.temporal_pooling = config.temporal_pooling
+        self._la_w1 = nn.Conv1d(config.hidden_size, int(config.hidden_size/2), 5, padding=2)
+        self._la_w2 = nn.Conv1d(config.hidden_size, int(config.hidden_size/2), 5, padding=2)
+        self._la_mlp = nn.Linear(config.hidden_size, config.hidden_size)
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        if self.temporal_pooling == 'mean':
+            return hidden_states.mean(dim=1)
+        if self.temporal_pooling == 'max':
+            return hidden_states.max(dim=1)
+        if self.temporal_pooling == 'concat':
+            _temp = hidden_states.reshape(hidden_states.shape[0], -1)
+            return torch.nn.functional.pad(_temp, (0, 2048 - _temp.shape[1]))
+        if self.temporal_pooling == 'topmax':
+            val, _ = torch.topk(hidden_states, k=5, dim=1)
+            return val.mean(dim=1)
+        if self.temporal_pooling == 'light_attention':
+            _temp = hidden_states.permute(0,2,1)
+            a = self._la_w1(_temp).softmax(dim=-1)
+            v = self._la_w2(_temp)
+            v_max = v.max(dim=-1).values
+            v_sum = (a * v).sum(dim=-1)
+            return self._la_mlp(torch.cat([v_max, v_sum], dim=1))
+
+        pooled_output = self.scalar_reweighting(hidden_states).squeeze(2)
+        pooled_output = self.dense(pooled_output)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class ProteinLSTMEncoder(nn.Module):
+
+    def __init__(self, config: ProteinLSTMConfig):
+        super().__init__()
+        forward_lstm = [ProteinLSTMLayer(config.input_size, config.hidden_size)]
+        reverse_lstm = [ProteinLSTMLayer(config.input_size, config.hidden_size)]
+        for _ in range(config.num_hidden_layers - 1):
+            forward_lstm.append(ProteinLSTMLayer(
+                config.hidden_size, config.hidden_size, config.hidden_dropout_prob))
+            reverse_lstm.append(ProteinLSTMLayer(
+                config.hidden_size, config.hidden_size, config.hidden_dropout_prob))
+        self.forward_lstm = nn.ModuleList(forward_lstm)
+        self.reverse_lstm = nn.ModuleList(reverse_lstm)
+        self.output_hidden_states = config.output_hidden_states
+
+    def forward(self, inputs, input_mask=None):
+        all_forward_pooled = ()
+        all_reverse_pooled = ()
+        all_hidden_states = (inputs,)
+        forward_output = inputs
+        for layer in self.forward_lstm:
+            forward_output, forward_pooled = layer(forward_output)
+            all_forward_pooled = all_forward_pooled + (forward_pooled[0],)
+            all_hidden_states = all_hidden_states + (forward_output,)
+
+        reversed_sequence = self.reverse_sequence(inputs, input_mask)
+        reverse_output = reversed_sequence
+        for layer in self.reverse_lstm:
+            reverse_output, reverse_pooled = layer(reverse_output)
+            all_reverse_pooled = all_reverse_pooled + (reverse_pooled[0],)
+            all_hidden_states = all_hidden_states + (reverse_output,)
+        reverse_output = self.reverse_sequence(reverse_output, input_mask)
+
+        output = torch.cat((forward_output, reverse_output), dim=2)
+
+        pooled = all_forward_pooled + all_reverse_pooled
+        pooled = torch.stack(pooled, 3).squeeze(0)
+        outputs = (output, pooled)
+        if self.output_hidden_states:
+            outputs = outputs + (all_hidden_states,)
+
+        return outputs  # sequence_embedding, pooled_embedding, (hidden_states)
+
+    def reverse_sequence(self, sequence, input_mask):
+        if input_mask is None:
+            idx = torch.arange(sequence.size(1) - 1, -1, -1)
+            reversed_sequence = sequence.index_select(1, idx, device=sequence.device)
+        else:
+            sequence_lengths = input_mask.sum(1)
+            reversed_sequence = []
+            for seq, seqlen in zip(sequence, sequence_lengths):
+                idx = torch.arange(seqlen - 1, -1, -1, device=seq.device)
+                seq = seq.index_select(0, idx)
+                seq = F.pad(seq, [0, 0, 0, sequence.size(1) - seqlen])
+                reversed_sequence.append(seq)
+            reversed_sequence = torch.stack(reversed_sequence, 0)
+        return reversed_sequence
+
+
+class ProteinLSTMAbstractModel(ProteinModel):
+
+    config_class = ProteinLSTMConfig
+    pretrained_model_archive_map = LSTM_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "lstm"
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+@registry.register_task_model('embed', 'lstm')
+class ProteinLSTMModel(ProteinLSTMAbstractModel):
+
+    def __init__(self, config: ProteinLSTMConfig):
+        super().__init__(config)
+        self.embed_matrix = nn.Embedding(config.vocab_size, config.input_size)
+        self.encoder = ProteinLSTMEncoder(config)
+        self.pooler = ProteinLSTMPooler(config)
+        self.output_hidden_states = config.output_hidden_states
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None):
+        if input_mask is None:
+            input_mask = torch.ones_like(input_ids)
+
+        # fp16 compatibility
+        embedding_output = self.embed_matrix(input_ids)
+        outputs = self.encoder(embedding_output, input_mask=input_mask)
+        sequence_output = outputs[0]
+        pooled_outputs = self.pooler(outputs[1])
+
+        outputs = (sequence_output, pooled_outputs) + outputs[2:]
+        return outputs  # sequence_output, pooled_output, (hidden_states)
+
+
+@registry.register_task_model('language_modeling', 'lstm')
+class ProteinLSTMForLM(ProteinLSTMAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.lstm = ProteinLSTMModel(config)
+        self.feedforward = nn.Linear(config.hidden_size, config.vocab_size)
+
+        self.init_weights()
+
+    def forward(self,
+                input_ids,
+                input_mask=None,
+                targets=None):
+
+        outputs = self.lstm(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+
+        forward_prediction, reverse_prediction = sequence_output.chunk(2, -1)
+        forward_prediction = F.pad(forward_prediction[:, :-1], [0, 0, 1, 0])
+        reverse_prediction = F.pad(reverse_prediction[:, 1:], [0, 0, 0, 1])
+        prediction_scores = \
+            self.feedforward(forward_prediction) + self.feedforward(reverse_prediction)
+        prediction_scores = prediction_scores.contiguous()
+
+        # add hidden states and if they are here
+        outputs = (prediction_scores,) + outputs[:2]
+
+        if targets is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
+            lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size), targets.view(-1))
+            outputs = (lm_loss,) + outputs
+
+        # (loss), prediction_scores, seq_relationship_score, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('fluorescence', 'lstm')
+@registry.register_task_model('stability', 'lstm')
+class ProteinLSTMForValuePrediction(ProteinLSTMAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.lstm = ProteinLSTMModel(config)
+        self.predict = ValuePredictionHead(config.hidden_size)
+        self.freeze_embedding = config.freeze_embedding
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+        if self.freeze_embedding:
+            self.lstm.train(False)
+
+        outputs = self.lstm(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('remote_homology', 'lstm')
+class ProteinLSTMForSequenceClassification(ProteinLSTMAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.lstm = ProteinLSTMModel(config)
+        self.classify = SequenceClassificationHead(
+            config.hidden_size, config.num_labels)
+        self.freeze_embedding = config.freeze_embedding
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+        if self.freeze_embedding:
+            self.lstm.train(False)
+
+        outputs = self.lstm(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('secondary_structure', 'lstm')
+class ProteinLSTMForSequenceToSequenceClassification(ProteinLSTMAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.lstm = ProteinLSTMModel(config)
+        self.classify = SequenceToSequenceClassificationHead(
+            config.hidden_size * 2, config.num_labels, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.lstm(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        amino_acid_class_scores = self.classify(sequence_output.contiguous())
+
+        # add hidden states and if they are here
+        outputs = (amino_acid_class_scores,) + outputs[2:]
+
+        if targets is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
+            classification_loss = loss_fct(
+                amino_acid_class_scores.view(-1, self.config.num_labels),
+                targets.view(-1))
+            outputs = (classification_loss,) + outputs
+
+        # (loss), prediction_scores, seq_relationship_score, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('contact_prediction', 'lstm')
+class ProteinLSTMForContactPrediction(ProteinLSTMAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.lstm = ProteinLSTMModel(config)
+        self.predict = PairwiseContactPredictionHead(config.hidden_size, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, protein_length, input_mask=None, targets=None):
+
+        outputs = self.lstm(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(sequence_output, protein_length, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs

tape/models/modeling_onehot.py

@@ -0,0 +1,155 @@
+import logging
+import typing
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .modeling_utils import ProteinConfig
+from .modeling_utils import ProteinModel
+from .modeling_utils import ValuePredictionHead
+from .modeling_utils import SequenceClassificationHead
+from .modeling_utils import SequenceToSequenceClassificationHead
+from .modeling_utils import PairwiseContactPredictionHead
+from ..registry import registry
+
+logger = logging.getLogger(__name__)
+
+
+class ProteinOneHotConfig(ProteinConfig):
+    pretrained_config_archive_map: typing.Dict[str, str] = {}
+
+    def __init__(self,
+                 vocab_size: int,
+                 initializer_range: float = 0.02,
+                 use_evolutionary: bool = False,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.use_evolutionary = use_evolutionary
+        self.initializer_range = initializer_range
+
+
+class ProteinOneHotAbstractModel(ProteinModel):
+
+    config_class = ProteinOneHotConfig
+    pretrained_model_archive_map: typing.Dict[str, str] = {}
+    base_model_prefix = "onehot"
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+class ProteinOneHotModel(ProteinOneHotAbstractModel):
+
+    def __init__(self, config: ProteinOneHotConfig):
+        super().__init__(config)
+        self.vocab_size = config.vocab_size
+
+        # Note: this exists *solely* for fp16 support
+        # There doesn't seem to be an easier way to check whether to use fp16 or fp32 training
+        buffer = torch.tensor([0.])
+        self.register_buffer('_buffer', buffer)
+
+    def forward(self, input_ids, input_mask=None):
+        if input_mask is None:
+            input_mask = torch.ones_like(input_ids)
+
+        sequence_output = F.one_hot(input_ids, num_classes=self.vocab_size)
+        # fp16 compatibility
+        sequence_output = sequence_output.type_as(self._buffer)
+        input_mask = input_mask.unsqueeze(2).type_as(sequence_output)
+        # just a bag-of-words for amino acids
+        pooled_outputs = (sequence_output * input_mask).sum(1) / input_mask.sum(1)
+
+        outputs = (sequence_output, pooled_outputs)
+        return outputs
+
+
+@registry.register_task_model('fluorescence', 'onehot')
+@registry.register_task_model('stability', 'onehot')
+class ProteinOneHotForValuePrediction(ProteinOneHotAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.onehot = ProteinOneHotModel(config)
+        self.predict = ValuePredictionHead(config.vocab_size)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.onehot(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('remote_homology', 'onehot')
+class ProteinOneHotForSequenceClassification(ProteinOneHotAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.onehot = ProteinOneHotModel(config)
+        self.classify = SequenceClassificationHead(config.vocab_size, config.num_labels)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.onehot(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('secondary_structure', 'onehot')
+class ProteinOneHotForSequenceToSequenceClassification(ProteinOneHotAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.onehot = ProteinOneHotModel(config)
+        self.classify = SequenceToSequenceClassificationHead(
+            config.vocab_size, config.num_labels, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.onehot(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(sequence_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('contact_prediction', 'onehot')
+class ProteinOneHotForContactPrediction(ProteinOneHotAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.onehot = ProteinOneHotModel(config)
+        self.predict = PairwiseContactPredictionHead(config.hidden_size, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, protein_length, input_mask=None, targets=None):
+
+        outputs = self.onehot(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(sequence_output, protein_length, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs

tape/models/modeling_resnet.py

@@ -0,0 +1,389 @@
+import typing
+import logging
+import torch
+import torch.nn as nn
+
+from .modeling_utils import ProteinConfig
+from .modeling_utils import ProteinModel
+from .modeling_utils import get_activation_fn
+from .modeling_utils import MLMHead
+from .modeling_utils import LayerNorm
+from .modeling_utils import ValuePredictionHead
+from .modeling_utils import SequenceClassificationHead
+from .modeling_utils import SequenceToSequenceClassificationHead
+from .modeling_utils import PairwiseContactPredictionHead
+from ..registry import registry
+
+logger = logging.getLogger(__name__)
+
+RESNET_PRETRAINED_CONFIG_ARCHIVE_MAP: typing.Dict[str, str] = {}
+RESNET_PRETRAINED_MODEL_ARCHIVE_MAP: typing.Dict[str, str] = {}
+
+
+class ProteinResNetConfig(ProteinConfig):
+    pretrained_config_archive_map = RESNET_PRETRAINED_CONFIG_ARCHIVE_MAP
+
+    def __init__(self,
+                 vocab_size: int = 30,
+                 hidden_size: int = 512,
+                 num_hidden_layers: int = 30,
+                 hidden_act: str = "gelu",
+                 hidden_dropout_prob: float = 0.1,
+                 initializer_range: float = 0.02,
+                 layer_norm_eps: float = 1e-12,
+                 temporal_pooling: str = 'attention',
+                 freeze_embedding: bool = False,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.num_hidden_layers = num_hidden_layers
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.temporal_pooling = temporal_pooling
+        self.freeze_embedding = freeze_embedding
+
+
+class MaskedConv1d(nn.Conv1d):
+
+    def forward(self, x, input_mask=None):
+        if input_mask is not None:
+            x = x * input_mask
+        return super().forward(x)
+
+
+class ProteinResNetLayerNorm(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.norm = LayerNorm(config.hidden_size)
+
+    def forward(self, x):
+        return self.norm(x.transpose(1, 2)).transpose(1, 2)
+
+
+class ProteinResNetBlock(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.conv1 = MaskedConv1d(
+            config.hidden_size, config.hidden_size, 3, padding=1, bias=False)
+        # self.bn1 = nn.BatchNorm1d(config.hidden_size)
+        self.bn1 = ProteinResNetLayerNorm(config)
+        self.conv2 = MaskedConv1d(
+            config.hidden_size, config.hidden_size, 3, padding=1, bias=False)
+        # self.bn2 = nn.BatchNorm1d(config.hidden_size)
+        self.bn2 = ProteinResNetLayerNorm(config)
+        self.activation_fn = get_activation_fn(config.hidden_act)
+
+    def forward(self, x, input_mask=None):
+        identity = x
+
+        out = self.conv1(x, input_mask)
+        out = self.bn1(out)
+        out = self.activation_fn(out)
+
+        out = self.conv2(out, input_mask)
+        out = self.bn2(out)
+
+        out += identity
+        out = self.activation_fn(out)
+
+        return out
+
+
+class ProteinResNetEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings.
+    """
+    def __init__(self, config):
+        super().__init__()
+        embed_dim = config.hidden_size
+        self.word_embeddings = nn.Embedding(config.vocab_size, embed_dim, padding_idx=0)
+        inverse_frequency = 1 / (10000 ** (torch.arange(0.0, embed_dim, 2.0) / embed_dim))
+        self.register_buffer('inverse_frequency', inverse_frequency)
+
+        self.layer_norm = LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, input_ids):
+        words_embeddings = self.word_embeddings(input_ids)
+
+        seq_length = input_ids.size(1)
+        position_ids = torch.arange(
+            seq_length - 1, -1, -1.0,
+            dtype=words_embeddings.dtype,
+            device=words_embeddings.device)
+        sinusoidal_input = torch.ger(position_ids, self.inverse_frequency)
+        position_embeddings = torch.cat([sinusoidal_input.sin(), sinusoidal_input.cos()], -1)
+        position_embeddings = position_embeddings.unsqueeze(0)
+
+        embeddings = words_embeddings + position_embeddings
+        embeddings = self.layer_norm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class ProteinResNetPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention_weights = nn.Linear(config.hidden_size, 1)
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+        self.temporal_pooling = config.temporal_pooling
+        self._la_w1 = nn.Conv1d(config.hidden_size, int(config.hidden_size/2), 5, padding=2)
+        self._la_w2 = nn.Conv1d(config.hidden_size, int(config.hidden_size/2), 5, padding=2)
+        self._la_mlp = nn.Linear(config.hidden_size, config.hidden_size)
+
+    def forward(self, hidden_states, mask=None):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        if self.temporal_pooling == 'mean':
+            return hidden_states.mean(dim=1)
+        if self.temporal_pooling == 'max':
+            return hidden_states.max(dim=1)
+        if self.temporal_pooling == 'concat':
+            _temp = hidden_states.reshape(hidden_states.shape[0], -1)
+            return torch.nn.functional.pad(_temp, (0, 2048 - _temp.shape[1]))
+        if self.temporal_pooling == 'meanmax':
+            _mean = hidden_states.mean(dim=1)
+            _max = hidden_states.max(dim=1)
+            return torch.cat([_mean, _max])
+        if self.temporal_pooling == 'topmax':
+            val, _ = torch.topk(hidden_states, k=5, dim=1)
+            return val.mean(dim=1)
+        if self.temporal_pooling == 'light_attention':
+            _temp = hidden_states.permute(0,2,1)
+            a = self._la_w1(_temp).softmax(dim=-1)
+            v = self._la_w2(_temp)
+            v_max = v.max(dim=-1).values
+            v_sum = (a * v).sum(dim=-1)
+            return self._la_mlp(torch.cat([v_max, v_sum], dim=1))
+
+        attention_scores = self.attention_weights(hidden_states)
+        if mask is not None:
+            attention_scores += -10000. * (1 - mask)
+        attention_weights = torch.softmax(attention_scores, -1)
+        weighted_mean_embedding = torch.matmul(
+            hidden_states.transpose(1, 2), attention_weights).squeeze(2)
+        pooled_output = self.dense(weighted_mean_embedding)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class ResNetEncoder(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.output_hidden_states = config.output_hidden_states
+        self.layer = nn.ModuleList(
+            [ProteinResNetBlock(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(self, hidden_states, input_mask=None):
+        all_hidden_states = ()
+        for layer_module in self.layer:
+            if self.output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            hidden_states = layer_module(hidden_states, input_mask)
+
+        if self.output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        outputs = (hidden_states,)
+        if self.output_hidden_states:
+            outputs = outputs + (all_hidden_states,)
+
+        return outputs
+
+
+class ProteinResNetAbstractModel(ProteinModel):
+    """ An abstract class to handle weights initialization and
+        a simple interface for dowloading and loading pretrained models.
+    """
+    config_class = ProteinResNetConfig
+    pretrained_model_archive_map = RESNET_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "resnet"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Conv1d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
+            if module.bias is not None:
+                module.bias.data.zero_()
+        # elif isinstance(module, ProteinResNetBlock):
+            # nn.init.constant_(module.bn2.weight, 0)
+
+
+@registry.register_task_model('embed', 'resnet')
+class ProteinResNetModel(ProteinResNetAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = ProteinResNetEmbeddings(config)
+        self.encoder = ResNetEncoder(config)
+        self.pooler = ProteinResNetPooler(config)
+
+        self.init_weights()
+
+    def forward(self,
+                input_ids,
+                input_mask=None):
+        if input_mask is not None and torch.any(input_mask != 1):
+            extended_input_mask = input_mask.unsqueeze(2)
+            # fp16 compatibility
+            extended_input_mask = extended_input_mask.to(
+                dtype=next(self.parameters()).dtype)
+        else:
+            extended_input_mask = None
+
+        embedding_output = self.embeddings(input_ids)
+        embedding_output = embedding_output.transpose(1, 2)
+        if extended_input_mask is not None:
+            extended_input_mask = extended_input_mask.transpose(1, 2)
+        encoder_outputs = self.encoder(embedding_output, extended_input_mask)
+        sequence_output = encoder_outputs[0]
+        sequence_output = sequence_output.transpose(1, 2).contiguous()
+        # sequence_output = encoder_outputs[0]
+        if extended_input_mask is not None:
+            extended_input_mask = extended_input_mask.transpose(1, 2)
+        pooled_output = self.pooler(sequence_output, extended_input_mask)
+
+        # add hidden_states and attentions if they are here
+        outputs = (sequence_output, pooled_output,) + encoder_outputs[1:]
+        return outputs  # sequence_output, pooled_output, (hidden_states)
+
+
+@registry.register_task_model('masked_language_modeling', 'resnet')
+class ProteinResNetForMaskedLM(ProteinResNetAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.resnet = ProteinResNetModel(config)
+        self.mlm = MLMHead(
+            config.hidden_size, config.vocab_size, config.hidden_act, config.layer_norm_eps,
+            ignore_index=-1)
+
+        self.init_weights()
+        self.tie_weights()
+
+    def tie_weights(self):
+        """ Make sure we are sharing the input and output embeddings.
+            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
+        """
+        self._tie_or_clone_weights(self.mlm.decoder,
+                                   self.resnet.embeddings.word_embeddings)
+
+    def forward(self,
+                input_ids,
+                input_mask=None,
+                targets=None):
+
+        outputs = self.resnet(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.mlm(sequence_output, targets) + outputs[:2]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('fluorescence', 'resnet')
+@registry.register_task_model('stability', 'resnet')
+class ProteinResNetForValuePrediction(ProteinResNetAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.resnet = ProteinResNetModel(config)
+        self.predict = ValuePredictionHead(config.hidden_size)
+        self.freeze_embedding = config.freeze_embedding
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+        if self.freeze_embedding:
+            self.resnet.train(False)
+
+        outputs = self.resnet(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('remote_homology', 'resnet')
+class ProteinResNetForSequenceClassification(ProteinResNetAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.resnet = ProteinResNetModel(config)
+        self.classify = SequenceClassificationHead(config.hidden_size, config.num_labels)
+        self.freeze_embedding = config.freeze_embedding
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+        if self.freeze_embedding:
+            self.resnet.train(False)
+
+        outputs = self.resnet(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('secondary_structure', 'resnet')
+class ProteinResNetForSequenceToSequenceClassification(ProteinResNetAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.resnet = ProteinResNetModel(config)
+        self.classify = SequenceToSequenceClassificationHead(
+            config.hidden_size, config.num_labels, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.resnet(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(sequence_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs
+
+
+@registry.register_task_model('contact_prediction', 'resnet')
+class ProteinResNetForContactPrediction(ProteinResNetAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.resnet = ProteinResNetModel(config)
+        self.predict = PairwiseContactPredictionHead(config.hidden_size, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, protein_length, input_mask=None, targets=None):
+
+        outputs = self.resnet(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(sequence_output, protein_length, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs

tape/models/modeling_trrosetta.py

@@ -0,0 +1,336 @@
+import torch
+import torch.nn as nn
+
+from ..registry import registry
+from .modeling_utils import ProteinConfig
+from .modeling_utils import ProteinModel
+
+URL_PREFIX = "https://s3.amazonaws.com/proteindata/pytorch-models/"
+TRROSETTA_PRETRAINED_MODEL_ARCHIVE_MAP = {
+    'xaa': URL_PREFIX + "trRosetta-xaa-pytorch_model.bin",
+    'xab': URL_PREFIX + "trRosetta-xab-pytorch_model.bin",
+    'xac': URL_PREFIX + "trRosetta-xac-pytorch_model.bin",
+    'xad': URL_PREFIX + "trRosetta-xad-pytorch_model.bin",
+    'xae': URL_PREFIX + "trRosetta-xae-pytorch_model.bin",
+}
+TRROSETTA_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    'xaa': URL_PREFIX + "trRosetta-xaa-config.json",
+    'xab': URL_PREFIX + "trRosetta-xab-config.json",
+    'xac': URL_PREFIX + "trRosetta-xac-config.json",
+    'xad': URL_PREFIX + "trRosetta-xad-config.json",
+    'xae': URL_PREFIX + "trRosetta-xae-config.json",
+}
+
+
+class TRRosettaConfig(ProteinConfig):
+
+    pretrained_config_archive_map = TRROSETTA_PRETRAINED_CONFIG_ARCHIVE_MAP
+
+    def __init__(self,
+                 num_features: int = 64,
+                 kernel_size: int = 3,
+                 num_layers: int = 61,
+                 dropout: float = 0.15,
+                 msa_cutoff: float = 0.8,
+                 penalty_coeff: float = 4.5,
+                 initializer_range: float = 0.02,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.num_features = num_features
+        self.kernel_size = kernel_size
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.msa_cutoff = msa_cutoff
+        self.penalty_coeff = penalty_coeff
+        self.initializer_range = initializer_range
+
+
+class MSAFeatureExtractor(nn.Module):
+
+    def __init__(self, config: TRRosettaConfig):
+        super().__init__()
+        self.msa_cutoff = config.msa_cutoff
+        self.penalty_coeff = config.penalty_coeff
+
+    def forward(self, msa1hot):
+        # Convert to float, then potentially back to half
+        # These transforms aren't well suited to half-precision
+        initial_type = msa1hot.dtype
+
+        msa1hot = msa1hot.float()
+        seqlen = msa1hot.size(2)
+
+        weights = self.reweight(msa1hot)
+        features_1d = self.extract_features_1d(msa1hot, weights)
+        features_2d = self.extract_features_2d(msa1hot, weights)
+
+        left = features_1d.unsqueeze(2).repeat(1, 1, seqlen, 1)
+        right = features_1d.unsqueeze(1).repeat(1, seqlen, 1, 1)
+        features = torch.cat((left, right, features_2d), -1)
+        features = features.type(initial_type)
+        features = features.permute(0, 3, 1, 2)
+        features = features.contiguous()
+        return features
+
+    def reweight(self, msa1hot, eps=1e-9):
+        # Reweight
+        seqlen = msa1hot.size(2)
+        id_min = seqlen * self.msa_cutoff
+        id_mtx = torch.stack([torch.tensordot(el, el, [[1, 2], [1, 2]]) for el in msa1hot], 0)
+        id_mask = id_mtx > id_min
+        weights = 1.0 / (id_mask.type_as(msa1hot).sum(-1) + eps)
+        return weights
+
+    def extract_features_1d(self, msa1hot, weights):
+        # 1D Features
+        f1d_seq = msa1hot[:, 0, :, :20]
+        batch_size = msa1hot.size(0)
+        seqlen = msa1hot.size(2)
+
+        # msa2pssm
+        beff = weights.sum()
+        f_i = (weights[:, :, None, None] * msa1hot).sum(1) / beff + 1e-9
+        h_i = (-f_i * f_i.log()).sum(2, keepdims=True)
+        f1d_pssm = torch.cat((f_i, h_i), dim=2)
+        f1d = torch.cat((f1d_seq, f1d_pssm), dim=2)
+        f1d = f1d.view(batch_size, seqlen, 42)
+        return f1d
+
+    def extract_features_2d(self, msa1hot, weights):
+        # 2D Features
+        batch_size = msa1hot.size(0)
+        num_alignments = msa1hot.size(1)
+        seqlen = msa1hot.size(2)
+        num_symbols = 21
+
+        if num_alignments == 1:
+            # No alignments, predict from sequence alone
+            f2d_dca = torch.zeros(
+                batch_size, seqlen, seqlen, 442,
+                dtype=torch.float,
+                device=msa1hot.device)
+            return f2d_dca
+
+        # compute fast_dca
+        # covariance
+        x = msa1hot.view(batch_size, num_alignments, seqlen * num_symbols)
+        num_points = weights.sum(1) - weights.mean(1).sqrt()
+        mean = (x * weights.unsqueeze(2)).sum(1, keepdims=True) / num_points[:, None, None]
+        x = (x - mean) * weights[:, :, None].sqrt()
+        cov = torch.matmul(x.transpose(-1, -2), x) / num_points[:, None, None]
+
+        # inverse covariance
+        reg = torch.eye(seqlen * num_symbols,
+                        device=weights.device,
+                        dtype=weights.dtype)[None]
+        reg = reg * self.penalty_coeff / weights.sum(1, keepdims=True).sqrt().unsqueeze(2)
+        cov_reg = cov + reg
+        inv_cov = torch.stack([torch.inverse(cr) for cr in cov_reg.unbind(0)], 0)
+
+        x1 = inv_cov.view(batch_size, seqlen, num_symbols, seqlen, num_symbols)
+        x2 = x1.permute(0, 1, 3, 2, 4)
+        features = x2.reshape(batch_size, seqlen, seqlen, num_symbols * num_symbols)
+
+        x3 = (x1[:, :, :-1, :, :-1] ** 2).sum((2, 4)).sqrt() * (
+            1 - torch.eye(seqlen, device=weights.device, dtype=weights.dtype)[None])
+        apc = x3.sum(1, keepdims=True) * x3.sum(2, keepdims=True) / x3.sum(
+            (1, 2), keepdims=True)
+        contacts = (x3 - apc) * (1 - torch.eye(
+            seqlen, device=x3.device, dtype=x3.dtype).unsqueeze(0))
+
+        f2d_dca = torch.cat([features, contacts[:, :, :, None]], axis=3)
+        return f2d_dca
+
+    @property
+    def feature_size(self) -> int:
+        return 526
+
+
+class DilatedResidualBlock(nn.Module):
+
+    def __init__(self, num_features: int, kernel_size: int, dilation: int, dropout: float):
+        super().__init__()
+        padding = self._get_padding(kernel_size, dilation)
+        self.conv1 = nn.Conv2d(
+            num_features, num_features, kernel_size, padding=padding, dilation=dilation)
+        self.norm1 = nn.InstanceNorm2d(num_features, affine=True, eps=1e-6)
+        self.actv1 = nn.ELU(inplace=True)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(
+            num_features, num_features, kernel_size, padding=padding, dilation=dilation)
+        self.norm2 = nn.InstanceNorm2d(num_features, affine=True, eps=1e-6)
+        self.actv2 = nn.ELU(inplace=True)
+        self.apply(self._init_weights)
+        nn.init.constant_(self.norm2.weight, 0)
+
+    def _get_padding(self, kernel_size: int, dilation: int) -> int:
+        return (kernel_size + (kernel_size - 1) * (dilation - 1) - 1) // 2
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+        # elif isinstance(module, DilatedResidualBlock):
+            # nn.init.constant_(module.norm2.weight, 0)
+
+    def forward(self, features):
+        shortcut = features
+        features = self.conv1(features)
+        features = self.norm1(features)
+        features = self.actv1(features)
+        features = self.dropout(features)
+        features = self.conv2(features)
+        features = self.norm2(features)
+        features = self.actv2(features + shortcut)
+        return features
+
+
+class TRRosettaAbstractModel(ProteinModel):
+
+    config_class = TRRosettaConfig
+    base_model_prefix = 'trrosetta'
+    pretrained_model_archive_map = TRROSETTA_PRETRAINED_MODEL_ARCHIVE_MAP
+
+    def __init__(self, config: TRRosettaConfig):
+        super().__init__(config)
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, DilatedResidualBlock):
+            nn.init.constant_(module.norm2.weight, 0)
+
+
+class TRRosettaPredictor(TRRosettaAbstractModel):
+
+    def __init__(self, config: TRRosettaConfig):
+        super().__init__(config)
+        layers = [
+            nn.Conv2d(526, config.num_features, 1),
+            nn.InstanceNorm2d(config.num_features, affine=True, eps=1e-6),
+            nn.ELU(),
+            nn.Dropout(config.dropout)]
+
+        dilation = 1
+        for _ in range(config.num_layers):
+            block = DilatedResidualBlock(
+                config.num_features, config.kernel_size, dilation, config.dropout)
+            layers.append(block)
+
+            dilation *= 2
+            if dilation > 16:
+                dilation = 1
+
+        self.resnet = nn.Sequential(*layers)
+        self.predict_theta = nn.Conv2d(config.num_features, 25, 1)
+        self.predict_phi = nn.Conv2d(config.num_features, 13, 1)
+        self.predict_dist = nn.Conv2d(config.num_features, 37, 1)
+        self.predict_bb = nn.Conv2d(config.num_features, 3, 1)
+        self.predict_omega = nn.Conv2d(config.num_features, 25, 1)
+
+        self.init_weights()
+
+    def init_weights(self):
+        self.apply(self._init_weights)
+        nn.init.constant_(self.predict_theta.weight, 0)
+        nn.init.constant_(self.predict_phi.weight, 0)
+        nn.init.constant_(self.predict_dist.weight, 0)
+        nn.init.constant_(self.predict_bb.weight, 0)
+        nn.init.constant_(self.predict_omega.weight, 0)
+
+    def forward(self,
+                features,
+                theta=None,
+                phi=None,
+                dist=None,
+                omega=None):
+        batch_size = features.size(0)
+        seqlen = features.size(2)
+        embedding = self.resnet(features)
+
+        # anglegrams for theta
+        logits_theta = self.predict_theta(embedding)
+
+        # anglegrams for phi
+        logits_phi = self.predict_phi(embedding)
+
+        # symmetrize
+        sym_embedding = 0.5 * (embedding + embedding.transpose(-1, -2))
+
+        # distograms
+        logits_dist = self.predict_dist(sym_embedding)
+
+        # beta-strand pairings (not used)
+        # logits_bb = self.predict_bb(sym_embedding)
+
+        # anglegrams for omega
+        logits_omega = self.predict_omega(sym_embedding)
+
+        logits_dist = logits_dist.permute(0, 2, 3, 1).contiguous()
+        logits_theta = logits_theta.permute(0, 2, 3, 1).contiguous()
+        logits_omega = logits_omega.permute(0, 2, 3, 1).contiguous()
+        logits_phi = logits_phi.permute(0, 2, 3, 1).contiguous()
+
+        probs = {}
+        probs['p_dist'] = nn.Softmax(-1)(logits_dist)
+        probs['p_theta'] = nn.Softmax(-1)(logits_theta)
+        probs['p_omega'] = nn.Softmax(-1)(logits_omega)
+        probs['p_phi'] = nn.Softmax(-1)(logits_phi)
+        outputs = (probs,)
+
+        metrics = {}
+        total_loss = 0
+
+        if dist is not None:
+            logits_dist = logits_dist.reshape(batch_size * seqlen * seqlen, 37)
+            loss_dist = nn.CrossEntropyLoss(ignore_index=-1)(logits_dist, dist.view(-1))
+            metrics['dist'] = loss_dist
+            total_loss += loss_dist
+        if theta is not None:
+            logits_theta = logits_theta.reshape(batch_size * seqlen * seqlen, 25)
+            loss_theta = nn.CrossEntropyLoss(ignore_index=0)(logits_theta, theta.view(-1))
+            metrics['theta'] = loss_theta
+            total_loss += loss_theta
+        if omega is not None:
+            logits_omega = logits_omega.reshape(batch_size * seqlen * seqlen, 25)
+            loss_omega = nn.CrossEntropyLoss(ignore_index=0)(logits_omega, omega.view(-1))
+            metrics['omega'] = loss_omega
+            total_loss += loss_omega
+        if phi is not None:
+            logits_phi = logits_phi.reshape(batch_size * seqlen * seqlen, 13)
+            loss_phi = nn.CrossEntropyLoss(ignore_index=0)(logits_phi, phi.view(-1))
+            metrics['phi'] = loss_phi
+            total_loss += loss_phi
+
+        if len(metrics) > 0:
+            outputs = ((total_loss, metrics),) + outputs
+
+        return outputs
+
+
+@registry.register_task_model('trrosetta', 'trrosetta')
+class TRRosetta(TRRosettaAbstractModel):
+
+    def __init__(self, config: TRRosettaConfig):
+        super().__init__(config)
+        self.extract_features = MSAFeatureExtractor(config)
+        self.trrosetta = TRRosettaPredictor(config)
+
+    def forward(self,
+                msa1hot,
+                theta=None,
+                phi=None,
+                dist=None,
+                omega=None):
+        features = self.extract_features(msa1hot)
+        return self.trrosetta(features, theta, phi, dist, omega)

tape/models/modeling_unirep.py

@@ -0,0 +1,270 @@
+import logging
+import typing
+import torch
+import torch.nn as nn
+from torch.nn.utils import weight_norm
+
+from .modeling_utils import ProteinConfig
+from .modeling_utils import ProteinModel
+from .modeling_utils import ValuePredictionHead
+from .modeling_utils import SequenceClassificationHead
+from .modeling_utils import SequenceToSequenceClassificationHead
+from .modeling_utils import PairwiseContactPredictionHead
+from ..registry import registry
+
+logger = logging.getLogger(__name__)
+
+
+URL_PREFIX = "https://s3.amazonaws.com/proteindata/pytorch-models/"
+UNIREP_PRETRAINED_CONFIG_ARCHIVE_MAP: typing.Dict[str, str] = {
+    'babbler-1900': URL_PREFIX + 'unirep-base-config.json'}
+UNIREP_PRETRAINED_MODEL_ARCHIVE_MAP: typing.Dict[str, str] = {
+    'babbler-1900': URL_PREFIX + 'unirep-base-pytorch_model.bin'}
+
+
+class UniRepConfig(ProteinConfig):
+    pretrained_config_archive_map = UNIREP_PRETRAINED_CONFIG_ARCHIVE_MAP
+
+    def __init__(self,
+                 vocab_size: int = 26,
+                 input_size: int = 10,
+                 hidden_size: int = 1900,
+                 hidden_dropout_prob: float = 0.1,
+                 layer_norm_eps: float = 1e-12,
+                 initializer_range: float = 0.02,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.layer_norm_eps = layer_norm_eps
+        self.initializer_range = initializer_range
+
+
+class mLSTMCell(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        project_size = config.hidden_size * 4
+        self.wmx = weight_norm(
+            nn.Linear(config.input_size, config.hidden_size, bias=False))
+        self.wmh = weight_norm(
+            nn.Linear(config.hidden_size, config.hidden_size, bias=False))
+        self.wx = weight_norm(
+            nn.Linear(config.input_size, project_size, bias=False))
+        self.wh = weight_norm(
+            nn.Linear(config.hidden_size, project_size, bias=True))
+
+    def forward(self, inputs, state):
+        h_prev, c_prev = state
+        m = self.wmx(inputs) * self.wmh(h_prev)
+        z = self.wx(inputs) + self.wh(m)
+        i, f, o, u = torch.chunk(z, 4, 1)
+        i = torch.sigmoid(i)
+        f = torch.sigmoid(f)
+        o = torch.sigmoid(o)
+        u = torch.tanh(u)
+        c = f * c_prev + i * u
+        h = o * torch.tanh(c)
+
+        return h, c
+
+
+class mLSTM(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.mlstm_cell = mLSTMCell(config)
+        self.hidden_size = config.hidden_size
+
+    def forward(self, inputs, state=None, mask=None):
+        batch_size = inputs.size(0)
+        seqlen = inputs.size(1)
+
+        if mask is None:
+            mask = torch.ones(batch_size, seqlen, 1, dtype=inputs.dtype, device=inputs.device)
+        elif mask.dim() == 2:
+            mask = mask.unsqueeze(2)
+
+        if state is None:
+            zeros = torch.zeros(batch_size, self.hidden_size,
+                                dtype=inputs.dtype, device=inputs.device)
+            state = (zeros, zeros)
+
+        steps = []
+        for seq in range(seqlen):
+            prev = state
+            seq_input = inputs[:, seq, :]
+            hx, cx = self.mlstm_cell(seq_input, state)
+            seqmask = mask[:, seq]
+            hx = seqmask * hx + (1 - seqmask) * prev[0]
+            cx = seqmask * cx + (1 - seqmask) * prev[1]
+            state = (hx, cx)
+            steps.append(hx)
+
+        return torch.stack(steps, 1), (hx, cx)
+
+
+class UniRepAbstractModel(ProteinModel):
+
+    config_class = UniRepConfig
+    pretrained_model_archive_map = UNIREP_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "unirep"
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+@registry.register_task_model('embed', 'unirep')
+class UniRepModel(UniRepAbstractModel):
+
+    def __init__(self, config: UniRepConfig):
+        super().__init__(config)
+        self.embed_matrix = nn.Embedding(config.vocab_size, config.input_size)
+        self.encoder = mLSTM(config)
+        self.output_hidden_states = config.output_hidden_states
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None):
+        if input_mask is None:
+            input_mask = torch.ones_like(input_ids)
+
+        # fp16 compatibility
+        input_mask = input_mask.to(dtype=next(self.parameters()).dtype)
+        embedding_output = self.embed_matrix(input_ids)
+
+        encoder_outputs = self.encoder(embedding_output, mask=input_mask)
+        sequence_output = encoder_outputs[0]
+        hidden_states = encoder_outputs[1]
+        pooled_outputs = torch.cat(hidden_states, 1)
+
+        outputs = (sequence_output, pooled_outputs)
+        return outputs
+
+
+@registry.register_task_model('language_modeling', 'unirep')
+class UniRepForLM(UniRepAbstractModel):
+    # TODO: Fix this for UniRep - UniRep changes the size of the targets
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.unirep = UniRepModel(config)
+        self.feedforward = nn.Linear(config.hidden_size, config.vocab_size - 1)
+
+        self.init_weights()
+
+    def forward(self,
+                input_ids,
+                input_mask=None,
+                targets=None):
+
+        outputs = self.unirep(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        prediction_scores = self.feedforward(sequence_output)
+
+        # add hidden states and if they are here
+        outputs = (prediction_scores,) + outputs[2:]
+
+        if targets is not None:
+            targets = targets[:, 1:]
+            prediction_scores = prediction_scores[:, :-1]
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
+            lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size), targets.view(-1))
+            outputs = (lm_loss,) + outputs
+
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('fluorescence', 'unirep')
+@registry.register_task_model('stability', 'unirep')
+class UniRepForValuePrediction(UniRepAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.unirep = UniRepModel(config)
+        self.predict = ValuePredictionHead(config.hidden_size * 2)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.unirep(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('remote_homology', 'unirep')
+class UniRepForSequenceClassification(UniRepAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.unirep = UniRepModel(config)
+        self.classify = SequenceClassificationHead(
+            config.hidden_size * 2, config.num_labels)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.unirep(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(pooled_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('secondary_structure', 'unirep')
+class UniRepForSequenceToSequenceClassification(UniRepAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.unirep = UniRepModel(config)
+        self.classify = SequenceToSequenceClassificationHead(
+            config.hidden_size, config.num_labels, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, input_mask=None, targets=None):
+
+        outputs = self.unirep(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.classify(sequence_output, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states)
+        return outputs
+
+
+@registry.register_task_model('contact_prediction', 'unirep')
+class UniRepForContactPrediction(UniRepAbstractModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.unirep = UniRepModel(config)
+        self.predict = PairwiseContactPredictionHead(config.hidden_size, ignore_index=-1)
+
+        self.init_weights()
+
+    def forward(self, input_ids, protein_length, input_mask=None, targets=None):
+
+        outputs = self.unirep(input_ids, input_mask=input_mask)
+
+        sequence_output, pooled_output = outputs[:2]
+        outputs = self.predict(sequence_output, protein_length, targets) + outputs[2:]
+        # (loss), prediction_scores, (hidden_states), (attentions)
+        return outputs

tape/models/modeling_utils.py

@@ -0,0 +1,887 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+# Modified by Roshan Rao
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Protein models."""
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+import typing
+import copy
+import json
+import logging
+import os
+from io import open
+import math
+from torch.nn.utils.weight_norm import weight_norm
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .file_utils import cached_path
+
+CONFIG_NAME = "config.json"
+WEIGHTS_NAME = "pytorch_model.bin"
+
+logger = logging.getLogger(__name__)
+
+
+class ProteinConfig(object):
+    """ Base class for all configuration classes.
+        Handles a few parameters common to all models' configurations as well as methods
+        for loading/downloading/saving configurations.
+
+        Class attributes (overridden by derived classes):
+            - ``pretrained_config_archive_map``: a python ``dict`` of with `short-cut-names`
+                (string) as keys and `url` (string) of associated pretrained model
+                configurations as values.
+
+        Parameters:
+            ``finetuning_task``: string, default `None`. Name of the task used to fine-tune
+                the model.
+            ``num_labels``: integer, default `2`. Number of classes to use when the model is
+                a classification model (sequences/tokens)
+            ``output_attentions``: boolean, default `False`. Should the model returns
+                attentions weights.
+            ``output_hidden_states``: string, default `False`. Should the model returns all
+                hidden-states.
+            ``torchscript``: string, default `False`. Is the model used with Torchscript.
+    """
+    pretrained_config_archive_map: typing.Dict[str, str] = {}
+
+    def __init__(self, **kwargs):
+        self.finetuning_task = kwargs.pop('finetuning_task', None)
+        self.num_labels = kwargs.pop('num_labels', 2)
+        self.output_attentions = kwargs.pop('output_attentions', False)
+        self.output_hidden_states = kwargs.pop('output_hidden_states', False)
+        self.torchscript = kwargs.pop('torchscript', False)
+
+    def save_pretrained(self, save_directory):
+        """ Save a configuration object to the directory `save_directory`, so that it
+            can be re-loaded using the :func:`~ProteinConfig.from_pretrained`
+            class method.
+        """
+        assert os.path.isdir(save_directory), "Saving path should be a directory where the " \
+                                              "model and configuration can be saved"
+
+        # If we save using the predefined names, we can load using `from_pretrained`
+        output_config_file = os.path.join(save_directory, CONFIG_NAME)
+
+        self.to_json_file(output_config_file)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        r""" Instantiate a :class:`~ProteinConfig`
+             (or a derived class) from a pre-trained model configuration.
+
+        Parameters:
+            pretrained_model_name_or_path: either:
+
+                - a string with the `shortcut name` of a pre-trained model configuration to
+                  load from cache or download, e.g.: ``bert-base-uncased``.
+                - a path to a `directory` containing a configuration file saved using the
+                  :func:`~ProteinConfig.save_pretrained` method,
+                  e.g.: ``./my_model_directory/``.
+                - a path or url to a saved configuration JSON `file`,
+                  e.g.: ``./my_model_directory/configuration.json``.
+
+            cache_dir: (`optional`) string:
+                Path to a directory in which a downloaded pre-trained model
+                configuration should be cached if the standard cache should not be used.
+
+            kwargs: (`optional`) dict:
+                key/value pairs with which to update the configuration object after loading.
+
+                - The values in kwargs of any keys which are configuration attributes will
+                  be used to override the loaded values.
+                - Behavior concerning key/value pairs whose keys are *not* configuration
+                  attributes is controlled by the `return_unused_kwargs` keyword parameter.
+
+            return_unused_kwargs: (`optional`) bool:
+
+                - If False, then this function returns just the final configuration object.
+                - If True, then this functions returns a tuple `(config, unused_kwargs)`
+                  where `unused_kwargs` is a dictionary consisting of the key/value pairs
+                  whose keys are not configuration attributes: ie the part of kwargs which
+                  has not been used to update `config` and is otherwise ignored.
+
+        Examples::
+
+            # We can't instantiate directly the base class `ProteinConfig` so let's
+              show the examples on a derived class: ProteinBertConfig
+            # Download configuration from S3 and cache.
+            config = ProteinBertConfig.from_pretrained('bert-base-uncased')
+            # E.g. config (or model) was saved using `save_pretrained('./test/saved_model/')`
+            config = ProteinBertConfig.from_pretrained('./test/saved_model/')
+            config = ProteinBertConfig.from_pretrained(
+                './test/saved_model/my_configuration.json')
+            config = ProteinBertConfig.from_pretrained(
+                'bert-base-uncased', output_attention=True, foo=False)
+            assert config.output_attention == True
+            config, unused_kwargs = BertConfig.from_pretrained(
+                'bert-base-uncased', output_attention=True,
+                foo=False, return_unused_kwargs=True)
+            assert config.output_attention == True
+            assert unused_kwargs == {'foo': False}
+
+        """
+        cache_dir = kwargs.pop('cache_dir', None)
+        return_unused_kwargs = kwargs.pop('return_unused_kwargs', False)
+
+        if pretrained_model_name_or_path in cls.pretrained_config_archive_map:
+            config_file = cls.pretrained_config_archive_map[pretrained_model_name_or_path]
+        elif os.path.isdir(pretrained_model_name_or_path):
+            config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
+        else:
+            config_file = pretrained_model_name_or_path
+        # redirect to the cache, if necessary
+        try:
+            resolved_config_file = cached_path(config_file, cache_dir=cache_dir)
+        except EnvironmentError:
+            if pretrained_model_name_or_path in cls.pretrained_config_archive_map:
+                logger.error("Couldn't reach server at '{}' to download pretrained model "
+                             "configuration file.".format(config_file))
+            else:
+                logger.error(
+                    "Model name '{}' was not found in model name list ({}). "
+                    "We assumed '{}' was a path or url but couldn't find any file "
+                    "associated to this path or url.".format(
+                        pretrained_model_name_or_path,
+                        ', '.join(cls.pretrained_config_archive_map.keys()),
+                        config_file))
+            return None
+        if resolved_config_file == config_file:
+            logger.info("loading configuration file {}".format(config_file))
+        else:
+            logger.info("loading configuration file {} from cache at {}".format(
+                config_file, resolved_config_file))
+
+        # Load config
+        config = cls.from_json_file(resolved_config_file)
+
+        # Update config with kwargs if needed
+        to_remove = []
+        for key, value in kwargs.items():
+            if hasattr(config, key):
+                setattr(config, key, value)
+                to_remove.append(key)
+        for key in to_remove:
+            kwargs.pop(key, None)
+
+        logger.info("Model config %s", config)
+        if return_unused_kwargs:
+            return config, kwargs
+        else:
+            return config
+
+    @classmethod
+    def from_dict(cls, json_object):
+        """Constructs a `Config` from a Python dictionary of parameters."""
+        config = cls(vocab_size_or_config_json_file=-1)
+        for key, value in json_object.items():
+            config.__dict__[key] = value
+        return config
+
+    @classmethod
+    def from_json_file(cls, json_file):
+        """Constructs a `BertConfig` from a json file of parameters."""
+        with open(json_file, "r", encoding='utf-8') as reader:
+            text = reader.read()
+        return cls.from_dict(json.loads(text))
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __repr__(self):
+        return str(self.to_json_string())
+
+    def to_dict(self):
+        """Serializes this instance to a Python dictionary."""
+        output = copy.deepcopy(self.__dict__)
+        return output
+
+    def to_json_string(self):
+        """Serializes this instance to a JSON string."""
+        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+
+    def to_json_file(self, json_file_path):
+        """ Save this instance to a json file."""
+        with open(json_file_path, "w", encoding='utf-8') as writer:
+            writer.write(self.to_json_string())
+
+
+class ProteinModel(nn.Module):
+    r""" Base class for all models.
+
+        :class:`~ProteinModel` takes care of storing the configuration of
+        the models and handles methods for loading/downloading/saving models as well as a
+        few methods commons to all models to (i) resize the input embeddings and (ii) prune
+        heads in the self-attention heads.
+
+        Class attributes (overridden by derived classes):
+            - ``config_class``: a class derived from :class:`~ProteinConfig`
+              to use as configuration class for this model architecture.
+            - ``pretrained_model_archive_map``: a python ``dict`` of with `short-cut-names`
+              (string) as keys and `url` (string) of associated pretrained weights as values.
+
+            - ``base_model_prefix``: a string indicating the attribute associated to the
+              base model in derived classes of the same architecture adding modules on top
+              of the base model.
+    """
+    config_class: typing.Type[ProteinConfig] = ProteinConfig
+    pretrained_model_archive_map: typing.Dict[str, str] = {}
+    base_model_prefix = ""
+
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__()
+        if not isinstance(config, ProteinConfig):
+            raise ValueError(
+                "Parameter config in `{}(config)` should be an instance of class "
+                "`ProteinConfig`. To create a model from a pretrained model use "
+                "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
+                    self.__class__.__name__, self.__class__.__name__
+                ))
+        # Save config in model
+        self.config = config
+
+    def _get_resized_embeddings(self, old_embeddings, new_num_tokens=None):
+        """ Build a resized Embedding Module from a provided token Embedding Module.
+            Increasing the size will add newly initialized vectors at the end
+            Reducing the size will remove vectors from the end
+
+        Args:
+            new_num_tokens: (`optional`) int
+                New number of tokens in the embedding matrix.
+                Increasing the size will add newly initialized vectors at the end
+                Reducing the size will remove vectors from the end
+                If not provided or None: return the provided token Embedding Module.
+        Return: ``torch.nn.Embeddings``
+            Pointer to the resized Embedding Module or the old Embedding Module if
+            new_num_tokens is None
+        """
+        if new_num_tokens is None:
+            return old_embeddings
+
+        old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
+        if old_num_tokens == new_num_tokens:
+            return old_embeddings
+
+        # Build new embeddings
+        new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim)
+        new_embeddings.to(old_embeddings.weight.device)
+
+        # initialize all new embeddings (in particular added tokens)
+        self.init_weights(new_embeddings)
+
+        # Copy word embeddings from the previous weights
+        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)
+        new_embeddings.weight.data[:num_tokens_to_copy, :] = \
+            old_embeddings.weight.data[:num_tokens_to_copy, :]
+
+        return new_embeddings
+
+    def _tie_or_clone_weights(self, first_module, second_module):
+        """ Tie or clone module weights depending of weither we are using TorchScript or not
+        """
+        if self.config.torchscript:
+            first_module.weight = nn.Parameter(second_module.weight.clone())
+        else:
+            first_module.weight = second_module.weight
+
+    def resize_token_embeddings(self, new_num_tokens=None):
+        """ Resize input token embeddings matrix of the model if
+            new_num_tokens != config.vocab_size. Take care of tying weights embeddings
+            afterwards if the model class has a `tie_weights()` method.
+
+        Arguments:
+
+            new_num_tokens: (`optional`) int:
+                New number of tokens in the embedding matrix. Increasing the size will add
+                newly initialized vectors at the end. Reducing the size will remove vectors
+                from the end. If not provided or None: does nothing and just returns a
+                pointer to the input tokens ``torch.nn.Embeddings`` Module of the model.
+
+        Return: ``torch.nn.Embeddings``
+            Pointer to the input tokens Embeddings Module of the model
+        """
+        base_model = getattr(self, self.base_model_prefix, self)  # get the base model if needed
+        model_embeds = base_model._resize_token_embeddings(new_num_tokens)
+        if new_num_tokens is None:
+            return model_embeds
+
+        # Update base model and current model config
+        self.config.vocab_size = new_num_tokens
+        base_model.vocab_size = new_num_tokens
+
+        # Tie weights again if needed
+        if hasattr(self, 'tie_weights'):
+            self.tie_weights()
+
+        return model_embeds
+
+    def init_weights(self):
+        """ Initialize and prunes weights if needed. """
+        # Initialize weights
+        self.apply(self._init_weights)
+
+        # Prune heads if needed
+        if getattr(self.config, 'pruned_heads', False):
+            self.prune_heads(self.config.pruned_heads)
+
+    def prune_heads(self, heads_to_prune):
+        """ Prunes heads of the base model.
+
+            Arguments:
+
+                heads_to_prune: dict with keys being selected layer indices (`int`) and
+                    associated values being the list of heads to prune in said layer
+                    (list of `int`).
+        """
+        base_model = getattr(self, self.base_model_prefix, self)  # get the base model if needed
+        base_model._prune_heads(heads_to_prune)
+
+    def save_pretrained(self, save_directory):
+        """ Save a model and its configuration file to a directory, so that it
+            can be re-loaded using the `:func:`~ProteinModel.from_pretrained`
+            ` class method.
+        """
+        assert os.path.isdir(save_directory), "Saving path should be a directory where "\
+                                              "the model and configuration can be saved"
+
+        # Only save the model it-self if we are using distributed training
+        model_to_save = self.module if hasattr(self, 'module') else self
+
+        # Save configuration file
+        model_to_save.config.save_pretrained(save_directory)
+
+        # If we save using the predefined names, we can load using `from_pretrained`
+        output_model_file = os.path.join(save_directory, WEIGHTS_NAME)
+
+        torch.save(model_to_save.state_dict(), output_model_file)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        r"""Instantiate a pretrained pytorch model from a pre-trained model configuration.
+
+        The model is set in evaluation mode by default using ``model.eval()``
+        (Dropout modules are deactivated)
+        To train the model, you should first set it back in training mode with ``model.train()``
+
+        The warning ``Weights from XXX not initialized from pretrained model`` means that
+        the weights of XXX do not come pre-trained with the rest of the model.
+        It is up to you to train those weights with a downstream fine-tuning task.
+
+        The warning ``Weights from XXX not used in YYY`` means that the layer XXX is not used
+        by YYY, therefore those weights are discarded.
+
+        Parameters:
+            pretrained_model_name_or_path: either:
+
+                - a string with the `shortcut name` of a pre-trained model to load from cache
+                  or download, e.g.: ``bert-base-uncased``.
+                - a path to a `directory` containing model weights saved using
+                  :func:`~ProteinModel.save_pretrained`,
+                  e.g.: ``./my_model_directory/``.
+
+            model_args: (`optional`) Sequence of positional arguments:
+                All remaning positional arguments will be passed to the underlying model's
+                ``__init__`` method
+
+            config: (`optional`) instance of a class derived from
+                :class:`~ProteinConfig`: Configuration for the model to
+                use instead of an automatically loaded configuation. Configuration can be
+                automatically loaded when:
+
+                - the model is a model provided by the library (loaded with the
+                  ``shortcut-name`` string of a pretrained model), or
+                - the model was saved using
+                  :func:`~ProteinModel.save_pretrained` and is reloaded
+                  by suppling the save directory.
+                - the model is loaded by suppling a local directory as
+                  ``pretrained_model_name_or_path`` and a configuration JSON file named
+                  `config.json` is found in the directory.
+
+            state_dict: (`optional`) dict:
+                an optional state dictionnary for the model to use instead of a state
+                dictionary loaded from saved weights file. This option can be used if you
+                want to create a model from a pretrained configuration but load your own
+                weights. In this case though, you should check if using
+                :func:`~ProteinModel.save_pretrained` and
+                :func:`~ProteinModel.from_pretrained` is not a
+                simpler option.
+
+            cache_dir: (`optional`) string:
+                Path to a directory in which a downloaded pre-trained model
+                configuration should be cached if the standard cache should not be used.
+
+            force_download: (`optional`) boolean, default False:
+                Force to (re-)download the model weights and configuration files and override
+                the cached versions if they exists.
+
+            resume_download: (`optional`) boolean, default False:
+                Do not delete incompletely recieved file. Attempt to resume the download if
+                such a file exists.
+
+            output_loading_info: (`optional`) boolean:
+                Set to ``True`` to also return a dictionnary containing missing keys,
+                unexpected keys and error messages.
+
+            kwargs: (`optional`) Remaining dictionary of keyword arguments:
+                Can be used to update the configuration object (after it being loaded) and
+                initiate the model. (e.g. ``output_attention=True``). Behave differently
+                depending on whether a `config` is provided or automatically loaded:
+
+                - If a configuration is provided with ``config``, ``**kwarg
+                  directly passed to the underlying model's ``__init__`` method (we assume
+                  all relevant updates to the configuration have already been done)
+                - If a configuration is not provided, ``kwargs`` will be first passed to the
+                  configuration class initialization function
+                  (:func:`~ProteinConfig.from_pretrained`). Each key of
+                  ``kwargs`` that corresponds to a configuration attribute will be used to
+                  override said attribute with the supplied ``kwargs`` value. Remaining keys
+                  that do not correspond to any configuration attribute will be passed to the
+                  underlying model's ``__init__`` function.
+
+        Examples::
+
+            # Download model and configuration from S3 and cache.
+            model = ProteinBertModel.from_pretrained('bert-base-uncased')
+            # E.g. model was saved using `save_pretrained('./test/saved_model/')`
+            model = ProteinBertModel.from_pretrained('./test/saved_model/')
+            # Update configuration during loading
+            model = ProteinBertModel.from_pretrained('bert-base-uncased', output_attention=True)
+            assert model.config.output_attention == True
+
+        """
+        config = kwargs.pop('config', None)
+        state_dict = kwargs.pop('state_dict', None)
+        cache_dir = kwargs.pop('cache_dir', None)
+        output_loading_info = kwargs.pop('output_loading_info', False)
+
+        force_download = kwargs.pop("force_download", False)
+        kwargs.pop("resume_download", False)
+
+        # Load config
+        if config is None:
+            config, model_kwargs = cls.config_class.from_pretrained(
+                pretrained_model_name_or_path, *model_args,
+                cache_dir=cache_dir, return_unused_kwargs=True,
+                # force_download=force_download,
+                # resume_download=resume_download,
+                **kwargs
+            )
+        else:
+            model_kwargs = kwargs
+
+        # Load model
+        if pretrained_model_name_or_path in cls.pretrained_model_archive_map:
+            archive_file = cls.pretrained_model_archive_map[pretrained_model_name_or_path]
+        elif os.path.isdir(pretrained_model_name_or_path):
+            archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+        else:
+            archive_file = pretrained_model_name_or_path
+        # redirect to the cache, if necessary
+        try:
+            resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir,
+                                                force_download=force_download)
+        except EnvironmentError:
+            if pretrained_model_name_or_path in cls.pretrained_model_archive_map:
+                logger.error(
+                    "Couldn't reach server at '{}' to download pretrained weights.".format(
+                        archive_file))
+            else:
+                logger.error(
+                    "Model name '{}' was not found in model name list ({}). "
+                    "We assumed '{}' was a path or url but couldn't find any file "
+                    "associated to this path or url.".format(
+                        pretrained_model_name_or_path,
+                        ', '.join(cls.pretrained_model_archive_map.keys()),
+                        archive_file))
+            return None
+        if resolved_archive_file == archive_file:
+            logger.info("loading weights file {}".format(archive_file))
+        else:
+            logger.info("loading weights file {} from cache at {}".format(
+                archive_file, resolved_archive_file))
+
+        # Instantiate model.
+        model = cls(config, *model_args, **model_kwargs)
+
+        if state_dict is None:
+            state_dict = torch.load(resolved_archive_file, map_location='cpu')
+
+        # Convert old format to new format if needed from a PyTorch state_dict
+        old_keys = []
+        new_keys = []
+        for key in state_dict.keys():
+            new_key = None
+            if 'gamma' in key:
+                new_key = key.replace('gamma', 'weight')
+            if 'beta' in key:
+                new_key = key.replace('beta', 'bias')
+            if new_key:
+                old_keys.append(key)
+                new_keys.append(new_key)
+        for old_key, new_key in zip(old_keys, new_keys):
+            state_dict[new_key] = state_dict.pop(old_key)
+
+        # Load from a PyTorch state_dict
+        missing_keys = []
+        unexpected_keys = []
+        error_msgs = []
+        # copy state_dict so _load_from_state_dict can modify it
+        metadata = getattr(state_dict, '_metadata', None)
+        state_dict = state_dict.copy()
+        if metadata is not None:
+            state_dict._metadata = metadata
+
+        def load(module, prefix=''):
+            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+            module._load_from_state_dict(
+                state_dict, prefix, local_metadata, True, missing_keys,
+                unexpected_keys, error_msgs)
+            for name, child in module._modules.items():
+                if child is not None:
+                    load(child, prefix + name + '.')
+
+        # Make sure we are able to load base models as well as derived models (with heads)
+        start_prefix = ''
+        model_to_load = model
+        if cls.base_model_prefix not in (None, ''):
+            if not hasattr(model, cls.base_model_prefix) and \
+                    any(s.startswith(cls.base_model_prefix) for s in state_dict.keys()):
+                start_prefix = cls.base_model_prefix + '.'
+            if hasattr(model, cls.base_model_prefix) and \
+                    not any(s.startswith(cls.base_model_prefix) for s in state_dict.keys()):
+                model_to_load = getattr(model, cls.base_model_prefix)
+
+        load(model_to_load, prefix=start_prefix)
+        if len(missing_keys) > 0:
+            logger.info("Weights of {} not initialized from pretrained model: {}".format(
+                model.__class__.__name__, missing_keys))
+        if len(unexpected_keys) > 0:
+            logger.info("Weights from pretrained model not used in {}: {}".format(
+                model.__class__.__name__, unexpected_keys))
+        if len(error_msgs) > 0:
+            raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
+                               model.__class__.__name__, "\n\t".join(error_msgs)))
+
+        if hasattr(model, 'tie_weights'):
+            model.tie_weights()  # make sure word embedding weights are still tied
+
+        # Set model in evaluation mode to desactivate DropOut modules by default
+        model.eval()
+
+        if output_loading_info:
+            loading_info = {
+                "missing_keys": missing_keys,
+                "unexpected_keys": unexpected_keys,
+                "error_msgs": error_msgs}
+            return model, loading_info
+
+        return model
+
+
+def prune_linear_layer(layer, index, dim=0):
+    """ Prune a linear layer (a model parameters) to keep only entries in index.
+        Return the pruned layer as a new layer with requires_grad=True.
+        Used to remove heads.
+    """
+    index = index.to(layer.weight.device)
+    W = layer.weight.index_select(dim, index).clone().detach()
+    if layer.bias is not None:
+        if dim == 1:
+            b = layer.bias.clone().detach()
+        else:
+            b = layer.bias[index].clone().detach()
+    new_size = list(layer.weight.size())
+    new_size[dim] = len(index)
+    new_layer = nn.Linear(
+        new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
+    new_layer.weight.requires_grad = False
+    new_layer.weight.copy_(W.contiguous())
+    new_layer.weight.requires_grad = True
+    if layer.bias is not None:
+        new_layer.bias.requires_grad = False
+        new_layer.bias.copy_(b.contiguous())
+        new_layer.bias.requires_grad = True
+    return new_layer
+
+
+def accuracy(logits, labels, ignore_index: int = -100):
+    with torch.no_grad():
+        valid_mask = (labels != ignore_index)
+        predictions = logits.float().argmax(-1)
+        correct = (predictions == labels) * valid_mask
+        return correct.sum().float() / valid_mask.sum().float()
+
+
+def gelu(x):
+    """Implementation of the gelu activation function.
+        For information: OpenAI GPT's gelu is slightly different
+            (and gives slightly different results):
+        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+        Also see https://arxiv.org/abs/1606.08415
+    """
+    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+def swish(x):
+    return x * torch.sigmoid(x)
+
+
+def get_activation_fn(name: str) -> typing.Callable:
+    if name == 'gelu':
+        return gelu
+    elif name == 'relu':
+        return torch.nn.functional.relu
+    elif name == 'swish':
+        return swish
+    else:
+        raise ValueError(f"Unrecognized activation fn: {name}")
+
+
+try:
+    from apex.normalization.fused_layer_norm import FusedLayerNorm as LayerNorm  # type: ignore
+except (ImportError, AttributeError):
+    logger.info("Better speed can be achieved with apex installed from "
+                "https://www.github.com/nvidia/apex .")
+
+    class LayerNorm(nn.Module):  # type: ignore
+        def __init__(self, hidden_size, eps=1e-12):
+            """Construct a layernorm module in the TF style (epsilon inside the square root).
+            """
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(hidden_size))
+            self.bias = nn.Parameter(torch.zeros(hidden_size))
+            self.variance_epsilon = eps
+
+        def forward(self, x):
+            u = x.mean(-1, keepdim=True)
+            s = (x - u).pow(2).mean(-1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.variance_epsilon)
+            return self.weight * x + self.bias
+
+
+class SimpleMLP(nn.Module):
+
+    def __init__(self,
+                 in_dim: int,
+                 hid_dim: int,
+                 out_dim: int,
+                 dropout: float = 0.):
+        super().__init__()
+        self.main = nn.Sequential(
+            weight_norm(nn.Linear(in_dim, hid_dim), dim=None),
+            nn.ReLU(),
+            nn.Dropout(dropout, inplace=True),
+            weight_norm(nn.Linear(hid_dim, out_dim), dim=None))
+
+    def forward(self, x):
+        return self.main(x)
+
+
+class SimpleConv(nn.Module):
+
+    def __init__(self,
+                 in_dim: int,
+                 hid_dim: int,
+                 out_dim: int,
+                 dropout: float = 0.):
+        super().__init__()
+        self.main = nn.Sequential(
+            nn.BatchNorm1d(in_dim),  # Added this
+            weight_norm(nn.Conv1d(in_dim, hid_dim, 5, padding=2), dim=None),
+            nn.ReLU(),
+            nn.Dropout(dropout, inplace=True),
+            weight_norm(nn.Conv1d(hid_dim, out_dim, 3, padding=1), dim=None))
+
+    def forward(self, x):
+        x = x.transpose(1, 2)
+        x = self.main(x)
+        x = x.transpose(1, 2).contiguous()
+        return x
+
+
+class Accuracy(nn.Module):
+
+    def __init__(self, ignore_index: int = -100):
+        super().__init__()
+        self.ignore_index = ignore_index
+
+    def forward(self, inputs, target):
+        return accuracy(inputs, target, self.ignore_index)
+
+
+class PredictionHeadTransform(nn.Module):
+
+    def __init__(self,
+                 hidden_size: int,
+                 hidden_act: typing.Union[str, typing.Callable] = 'gelu',
+                 layer_norm_eps: float = 1e-12):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        if isinstance(hidden_act, str):
+            self.transform_act_fn = get_activation_fn(hidden_act)
+        else:
+            self.transform_act_fn = hidden_act
+        self.LayerNorm = LayerNorm(hidden_size, eps=layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+class MLMHead(nn.Module):
+
+    def __init__(self,
+                 hidden_size: int,
+                 vocab_size: int,
+                 hidden_act: typing.Union[str, typing.Callable] = 'gelu',
+                 layer_norm_eps: float = 1e-12,
+                 ignore_index: int = -100):
+        super().__init__()
+        self.transform = PredictionHeadTransform(hidden_size, hidden_act, layer_norm_eps)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(hidden_size, vocab_size, bias=False)
+        self.bias = nn.Parameter(data=torch.zeros(vocab_size))  # type: ignore
+        self.vocab_size = vocab_size
+        self._ignore_index = ignore_index
+
+    def forward(self, hidden_states, targets=None):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states) + self.bias
+        outputs = (hidden_states,)
+        if targets is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=self._ignore_index)
+            masked_lm_loss = loss_fct(
+                hidden_states.reshape(-1, self.vocab_size), targets.reshape(-1))
+            metrics = {'perplexity': torch.exp(masked_lm_loss)}
+            loss_and_metrics = (masked_lm_loss, metrics)
+            outputs = (loss_and_metrics,) + outputs
+        return outputs  # (loss), prediction_scores
+
+
+class ValuePredictionHead(nn.Module):
+    def __init__(self, hidden_size: int, dropout: float = 0.):
+        super().__init__()
+        self.value_prediction = SimpleMLP(hidden_size, 512, 1, dropout)
+
+    def forward(self, pooled_output, targets=None):
+        value_pred = self.value_prediction(pooled_output)
+        outputs = (value_pred,)
+
+        if targets is not None:
+            loss_fct = nn.MSELoss()
+            value_pred_loss = loss_fct(value_pred, targets)
+            outputs = (value_pred_loss,) + outputs
+        return outputs  # (loss), value_prediction
+
+
+class SequenceClassificationHead(nn.Module):
+    def __init__(self, hidden_size: int, num_labels: int):
+        super().__init__()
+        self.classify = SimpleMLP(hidden_size, 512, num_labels)
+
+    def forward(self, pooled_output, targets=None):
+        logits = self.classify(pooled_output)
+        outputs = (logits,)
+
+        if targets is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            classification_loss = loss_fct(logits, targets)
+            metrics = {'accuracy': accuracy(logits, targets)}
+            loss_and_metrics = (classification_loss, metrics)
+            outputs = (loss_and_metrics,) + outputs
+
+        return outputs  # (loss), logits
+
+
+class SequenceToSequenceClassificationHead(nn.Module):
+
+    def __init__(self,
+                 hidden_size: int,
+                 num_labels: int,
+                 ignore_index: int = -100):
+        super().__init__()
+        self.classify = SimpleConv(
+            hidden_size, 512, num_labels)
+        self.num_labels = num_labels
+        self._ignore_index = ignore_index
+
+    def forward(self, sequence_output, targets=None):
+        sequence_logits = self.classify(sequence_output)
+        outputs = (sequence_logits,)
+        if targets is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=self._ignore_index)
+            classification_loss = loss_fct(
+                sequence_logits.view(-1, self.num_labels), targets.view(-1))
+            acc_fct = Accuracy(ignore_index=self._ignore_index)
+            metrics = {'accuracy':
+                       acc_fct(sequence_logits.view(-1, self.num_labels), targets.view(-1))}
+            loss_and_metrics = (classification_loss, metrics)
+            outputs = (loss_and_metrics,) + outputs
+        return outputs  # (loss), sequence_logits
+
+
+class PairwiseContactPredictionHead(nn.Module):
+
+    def __init__(self, hidden_size: int, ignore_index=-100):
+        super().__init__()
+        self.predict = nn.Sequential(
+            nn.Dropout(), nn.Linear(2 * hidden_size, 2))
+        self._ignore_index = ignore_index
+
+    def forward(self, inputs, sequence_lengths, targets=None):
+        prod = inputs[:, :, None, :] * inputs[:, None, :, :]
+        diff = inputs[:, :, None, :] - inputs[:, None, :, :]
+        pairwise_features = torch.cat((prod, diff), -1)
+        prediction = self.predict(pairwise_features)
+        prediction = (prediction + prediction.transpose(1, 2)) / 2
+        prediction = prediction[:, 1:-1, 1:-1].contiguous()  # remove start/stop tokens
+        outputs = (prediction,)
+
+        if targets is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=self._ignore_index)
+            contact_loss = loss_fct(
+                prediction.view(-1, 2), targets.view(-1))
+            metrics = {'precision_at_l5':
+                       self.compute_precision_at_l5(sequence_lengths, prediction, targets)}
+            loss_and_metrics = (contact_loss, metrics)
+            outputs = (loss_and_metrics,) + outputs
+
+        return outputs
+
+    def compute_precision_at_l5(self, sequence_lengths, prediction, labels):
+        with torch.no_grad():
+            valid_mask = labels != self._ignore_index
+            seqpos = torch.arange(valid_mask.size(1), device=sequence_lengths.device)
+            x_ind, y_ind = torch.meshgrid(seqpos, seqpos)
+            valid_mask &= ((y_ind - x_ind) >= 6).unsqueeze(0)
+            probs = F.softmax(prediction, 3)[:, :, :, 1]
+            valid_mask = valid_mask.type_as(probs)
+            correct = 0
+            total = 0
+            for length, prob, label, mask in zip(sequence_lengths, probs, labels, valid_mask):
+                masked_prob = (prob * mask).view(-1)
+                most_likely = masked_prob.topk(length // 5, sorted=False)
+                selected = label.view(-1).gather(0, most_likely.indices)
+                correct += selected.sum().float()
+                total += selected.numel()
+            return correct / total

tape/optimization.py

@@ -0,0 +1,209 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Modifications by Roshan Rao
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch optimization for BERT model."""
+
+import logging
+import math
+
+import torch
+from torch.optim import Optimizer  # type: ignore
+from torch.optim.lr_scheduler import LambdaLR
+
+logger = logging.getLogger(__name__)
+
+
+class ConstantLRSchedule(LambdaLR):
+    """ Constant learning rate schedule.
+    """
+    def __init__(self, optimizer, last_epoch=-1):
+        super(ConstantLRSchedule, self).__init__(
+            optimizer, lambda _: 1.0, last_epoch=last_epoch)
+
+
+class WarmupConstantSchedule(LambdaLR):
+    """ Linear warmup and then constant.
+        Linearly increases learning rate schedule from 0 to 1 over `warmup_steps`
+        training steps. Keeps learning rate schedule equal to 1. after warmup_steps.
+    """
+    def __init__(self, optimizer, warmup_steps, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        super(WarmupConstantSchedule, self).__init__(
+            optimizer, self.lr_lambda, last_epoch=last_epoch)
+
+    def lr_lambda(self, step):
+        if step < self.warmup_steps:
+            return float(step) / float(max(1.0, self.warmup_steps))
+        return 1.
+
+
+class WarmupLinearSchedule(LambdaLR):
+    """ Linear warmup and then linear decay.
+        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
+        Linearly decreases learning rate from 1. to 0. over remaining `t_total - warmup_steps`
+        steps.
+    """
+    def __init__(self, optimizer, warmup_steps, t_total, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        self.t_total = t_total
+        super(WarmupLinearSchedule, self).__init__(
+            optimizer, self.lr_lambda, last_epoch=last_epoch)
+
+    def lr_lambda(self, step):
+        if step < self.warmup_steps:
+            return float(step) / float(max(1, self.warmup_steps))
+        return max(0.0, float(self.t_total - step) / float(
+            max(1.0, self.t_total - self.warmup_steps)))
+
+
+class WarmupCosineSchedule(LambdaLR):
+    """ Linear warmup and then cosine decay.
+        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
+        Decreases learning rate from 1. to 0. over remaining `t_total - warmup_steps` steps
+        following a cosine curve. If `cycles` (default=0.5) is different from default, learning
+        rate follows cosine function after warmup.
+    """
+    def __init__(self, optimizer, warmup_steps, t_total, cycles=.5, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        self.t_total = t_total
+        self.cycles = cycles
+        super(WarmupCosineSchedule, self).__init__(
+            optimizer, self.lr_lambda, last_epoch=last_epoch)
+
+    def lr_lambda(self, step):
+        if step < self.warmup_steps:
+            return float(step) / float(max(1.0, self.warmup_steps))
+        # progress after warmup
+        progress = float(step - self.warmup_steps) / float(
+            max(1, self.t_total - self.warmup_steps))
+        return max(0.0, 0.5 * (1. + math.cos(math.pi * float(self.cycles) * 2.0 * progress)))
+
+
+class WarmupCosineWithHardRestartsSchedule(LambdaLR):
+    """ Linear warmup and then cosine cycles with hard restarts.
+        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
+        If `cycles` (default=1.) is different from default, learning rate follows `cycles` times
+        a cosine decaying learning rate (with hard restarts).
+    """
+    def __init__(self, optimizer, warmup_steps, t_total, cycles=1., last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        self.t_total = t_total
+        self.cycles = cycles
+        super(WarmupCosineWithHardRestartsSchedule, self).__init__(
+            optimizer, self.lr_lambda, last_epoch=last_epoch)
+
+    def lr_lambda(self, step):
+        if step < self.warmup_steps:
+            return float(step) / float(max(1, self.warmup_steps))
+        # progress after warmup
+        progress = float(step - self.warmup_steps) / float(
+            max(1, self.t_total - self.warmup_steps))
+        if progress >= 1.0:
+            return 0.0
+        return max(0.0, 0.5 * (1. + math.cos(
+            math.pi * ((float(self.cycles) * progress) % 1.0))))
+
+
+class AdamW(Optimizer):
+    """ Implements Adam algorithm with weight decay fix.
+
+    Parameters:
+        lr (float): learning rate. Default 1e-3.
+        betas (tuple of 2 floats): Adams beta parameters (b1, b2). Default: (0.9, 0.999)
+        eps (float): Adams epsilon. Default: 1e-6
+        weight_decay (float): Weight decay. Default: 0.0
+        correct_bias (bool): can be set to False to avoid correcting bias in Adam
+            (e.g. like in Bert TF repository). Default True.
+    """
+    def __init__(self,
+                 params,
+                 lr=1e-3,
+                 betas=(0.9, 0.999),
+                 eps=1e-6,
+                 weight_decay=0.0,
+                 correct_bias=True):
+        if lr < 0.0:
+            raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)")
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {} - should be >= 0.0".format(eps))
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
+                        correct_bias=correct_bias)
+        super(AdamW, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Adam does not support sparse gradients, '
+                                       'please consider SparseAdam instead')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                # Decay the first and second moment running average coefficient
+                # In-place operations to update the averages at the same time
+                exp_avg.mul_(beta1).add_(1.0 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1.0 - beta2, grad, grad)
+                denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                step_size = group['lr']
+                if group['correct_bias']:  # No bias correction for Bert
+                    bias_correction1 = 1.0 - beta1 ** state['step']
+                    bias_correction2 = 1.0 - beta2 ** state['step']
+                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
+
+                p.data.addcdiv_(-step_size, exp_avg, denom)
+
+                # Just adding the square of the weights to the loss function is *not*
+                # the correct way of using L2 regularization/weight decay with Adam,
+                # since that will interact with the m and v parameters in strange ways.
+                #
+                # Instead we want to decay the weights in a manner that doesn't interact
+                # with the m/v parameters. This is equivalent to adding the square
+                # of the weights to the loss with plain (non-momentum) SGD.
+                # Add weight decay at the end (fixed version)
+                if group['weight_decay'] > 0.0:
+                    p.data.add_(-group['lr'] * group['weight_decay'], p.data)
+
+        return loss

tape/registry.py

@@ -0,0 +1,254 @@
+from typing import Dict, Type, Callable, Optional, Union
+from torch.utils.data import Dataset
+from .models.modeling_utils import ProteinModel
+from pathlib import Path
+
+PathType = Union[str, Path]
+
+
+def convert_model_args(model_args):
+    d = {}
+    for e in model_args:
+        k, v = e.split("=")
+        try:
+            v = int(v)
+        except:
+            try:
+                v = float(v)
+            except:
+                v = str(v)
+        d[k] = v
+    return d
+
+
+class TAPETaskSpec:
+    """
+    Attributes
+    ----------
+    name (str):
+        The name of the TAPE task
+    dataset (Type[Dataset]):
+        The dataset used in the TAPE task
+    num_labels (int):
+        number of labels used if this is a classification task
+    models (Dict[str, ProteinModel]):
+        The set of models that can be used for this task. Default: {}.
+    """
+
+    def __init__(self,
+                 name: str,
+                 dataset: Type[Dataset],
+                 num_labels: int = -1,
+                 models: Optional[Dict[str, Type[ProteinModel]]] = None):
+        self.name = name
+        self.dataset = dataset
+        self.num_labels = num_labels
+        self.models = models if models is not None else {}
+
+    def register_model(self, model_name: str, model_cls: Optional[Type[ProteinModel]] = None):
+        if model_cls is not None:
+            if model_name in self.models:
+                raise KeyError(
+                    f"A model with name '{model_name}' is already registered for this task")
+            self.models[model_name] = model_cls
+            return model_cls
+        else:
+            return lambda model_cls: self.register_model(model_name, model_cls)
+
+    def get_model(self, model_name: str) -> Type[ProteinModel]:
+        return self.models[model_name]
+
+
+class Registry:
+    r"""Class for registry object which acts as the
+    central repository for TAPE."""
+
+    task_name_mapping: Dict[str, TAPETaskSpec] = {}
+    metric_name_mapping: Dict[str, Callable] = {}
+
+    @classmethod
+    def register_task(cls,
+                      task_name: str,
+                      num_labels: int = -1,
+                      dataset: Optional[Type[Dataset]] = None,
+                      models: Optional[Dict[str, Type[ProteinModel]]] = None):
+        """ Register a a new TAPE task. This creates a new TAPETaskSpec.
+
+        Args:
+
+            task_name (str): The name of the TAPE task.
+            num_labels (int): Number of labels used if this is a classification task. If this
+                is not a classification task, simply leave the default as -1.
+            dataset (Type[Dataset]): The dataset used in the TAPE task.
+            models (Optional[Dict[str, ProteinModel]]): The set of models that can be used for
+                this task. If you do not pass this argument, you can register models to the task
+                later by using `registry.register_task_model`. Default: {}.
+
+        Examples:
+
+        There are two ways of registering a new task. First, one can define the task by simply
+        declaring all the components, and then calling the register method, like so:
+
+            class SecondaryStructureDataset(Dataset):
+                ...
+
+            class ProteinBertForSequenceToSequenceClassification():
+                ...
+
+            registry.register_task(
+                'secondary_structure', 3, SecondaryStructureDataset,
+                {'transformer': ProteinBertForSequenceToSequenceClassification})
+
+        This will register a new task, 'secondary_structure', with a single model. More models
+        can be added with `registry.register_task_model`. Alternatively, this can be used as a
+        decorator:
+
+            @registry.regsiter_task('secondary_structure', 3)
+            class SecondaryStructureDataset(Dataset):
+                ...
+
+            @registry.register_task_model('secondary_structure', 'transformer')
+            class ProteinBertForSequenceToSequenceClassification():
+                ...
+
+        These two pieces of code are exactly equivalent, in terms of the resulting registry
+        state.
+
+        """
+        if dataset is not None:
+            if models is None:
+                models = {}
+            task_spec = TAPETaskSpec(task_name, dataset, num_labels, models)
+            return cls.register_task_spec(task_name, task_spec).dataset
+        else:
+            return lambda dataset: cls.register_task(task_name, num_labels, dataset, models)
+
+    @classmethod
+    def register_task_spec(cls, task_name: str, task_spec: Optional[TAPETaskSpec] = None):
+        """ Registers a task_spec directly. If you find it easier to actually create a
+            TAPETaskSpec manually, and then register it, feel free to use this method,
+            but otherwise it is likely easier to use `registry.register_task`.
+        """
+        if task_spec is not None:
+            if task_name in cls.task_name_mapping:
+                raise KeyError(f"A task with name '{task_name}' is already registered")
+            cls.task_name_mapping[task_name] = task_spec
+            return task_spec
+        else:
+            return lambda task_spec: cls.register_task_spec(task_name, task_spec)
+
+    @classmethod
+    def register_task_model(cls,
+                            task_name: str,
+                            model_name: str,
+                            model_cls: Optional[Type[ProteinModel]] = None):
+        r"""Register a specific model to a task with the provided model name.
+            The task must already be in the registry - you cannot register a
+            model to an unregistered task.
+
+        Args:
+            task_name (str): Name of task to which to register the model.
+            model_name (str): Name of model to use when registering task, this
+                is the name that you will use to refer to the model on the
+                command line.
+            model_cls (Type[ProteinModel]): The model to register.
+
+        Examples:
+
+        As with `registry.register_task`, this can both be used as a regular
+        python function, and as a decorator. For example this:
+
+            class ProteinBertForSequenceToSequenceClassification():
+                ...
+            registry.register_task_model(
+                'secondary_structure', 'transformer',
+                ProteinBertForSequenceToSequenceClassification)
+
+        and as a decorator:
+
+            @registry.register_task_model('secondary_structure', 'transformer')
+            class ProteinBertForSequenceToSequenceClassification():
+                ...
+
+        are both equivalent.
+        """
+        if task_name not in cls.task_name_mapping:
+            raise KeyError(
+                f"Tried to register a task model for an unregistered task: {task_name}. "
+                f"Make sure to register the task {task_name} first.")
+        return cls.task_name_mapping[task_name].register_model(model_name, model_cls)
+
+    @classmethod
+    def register_metric(cls, name: str) -> Callable[[Callable], Callable]:
+        r"""Register a metric to registry with key 'name'
+
+        Args:
+            name: Key with which the metric will be registered.
+
+        Usage::
+            from tape.registry import registry
+
+            @registry.register_metric('mse')
+            def mean_squred_error(inputs, outputs):
+                ...
+        """
+
+        def wrap(fn: Callable) -> Callable:
+            assert callable(fn), "All metrics must be callable"
+            cls.metric_name_mapping[name] = fn
+            return fn
+
+        return wrap
+
+    @classmethod
+    def get_task_spec(cls, name: str) -> TAPETaskSpec:
+        return cls.task_name_mapping[name]
+
+    @classmethod
+    def get_metric(cls, name: str) -> Callable:
+        return cls.metric_name_mapping[name]
+
+    @classmethod
+    def get_task_model(cls,
+                       model_name: str,
+                       task_name: str,
+                       config_file: Optional[PathType] = None,
+                       load_dir: Optional[PathType] = None,
+                       model_args = None) -> ProteinModel:
+        """ Create a TAPE task model, either from scratch or from a pretrained model.
+            This is mostly a helper function that evaluates the if statements in a
+            sensible order if you pass all three of the arguments.
+        Args:
+            model_name (str): Which type of model to create (e.g. transformer, unirep, ...)
+            task_name (str): The TAPE task for which to create a model
+            config_file (str, optional): A json config file that specifies hyperparameters
+            load_dir (str, optional): A save directory for a pretrained model
+        Returns:
+            model (ProteinModel): A TAPE task model
+        """
+        task_spec = registry.get_task_spec(task_name)
+        model_cls = task_spec.get_model(model_name)
+
+        if load_dir is not None:
+            model = model_cls.from_pretrained(load_dir, num_labels=task_spec.num_labels)
+        else:
+            config_class = model_cls.config_class
+            if config_file is not None:
+                config = config_class.from_json_file(config_file)
+            else:
+                config = config_class()
+            
+            if model_args:
+                model_args = convert_model_args(model_args)
+                for k,v in model_args.items():
+                    if k in config.__dict__ and type(config.__dict__[k])==type(v):
+                        setattr(config, k, v)
+                    else:
+                        raise ValueError(f"model arg {k} not in config or of the same type as default")
+
+            config.num_labels = task_spec.num_labels
+            model = model_cls(config)
+        return model
+
+
+registry = Registry()

tape/tokenizers.py

@@ -0,0 +1,174 @@
+from typing import List
+import logging
+from collections import OrderedDict
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+IUPAC_CODES = OrderedDict([
+    ('Ala', 'A'),
+    ('Asx', 'B'),
+    ('Cys', 'C'),
+    ('Asp', 'D'),
+    ('Glu', 'E'),
+    ('Phe', 'F'),
+    ('Gly', 'G'),
+    ('His', 'H'),
+    ('Ile', 'I'),
+    ('Lys', 'K'),
+    ('Leu', 'L'),
+    ('Met', 'M'),
+    ('Asn', 'N'),
+    ('Pro', 'P'),
+    ('Gln', 'Q'),
+    ('Arg', 'R'),
+    ('Ser', 'S'),
+    ('Thr', 'T'),
+    ('Sec', 'U'),
+    ('Val', 'V'),
+    ('Trp', 'W'),
+    ('Xaa', 'X'),
+    ('Tyr', 'Y'),
+    ('Glx', 'Z')])
+
+IUPAC_VOCAB = OrderedDict([
+    ("<pad>", 0),
+    ("<mask>", 1),
+    ("<cls>", 2),
+    ("<sep>", 3),
+    ("<unk>", 4),
+    ("A", 5),
+    ("B", 6),
+    ("C", 7),
+    ("D", 8),
+    ("E", 9),
+    ("F", 10),
+    ("G", 11),
+    ("H", 12),
+    ("I", 13),
+    ("K", 14),
+    ("L", 15),
+    ("M", 16),
+    ("N", 17),
+    ("O", 18),
+    ("P", 19),
+    ("Q", 20),
+    ("R", 21),
+    ("S", 22),
+    ("T", 23),
+    ("U", 24),
+    ("V", 25),
+    ("W", 26),
+    ("X", 27),
+    ("Y", 28),
+    ("Z", 29)])
+
+UNIREP_VOCAB = OrderedDict([
+    ("<pad>", 0),
+    ("M", 1),
+    ("R", 2),
+    ("H", 3),
+    ("K", 4),
+    ("D", 5),
+    ("E", 6),
+    ("S", 7),
+    ("T", 8),
+    ("N", 9),
+    ("Q", 10),
+    ("C", 11),
+    ("U", 12),
+    ("G", 13),
+    ("P", 14),
+    ("A", 15),
+    ("V", 16),
+    ("I", 17),
+    ("F", 18),
+    ("Y", 19),
+    ("W", 20),
+    ("L", 21),
+    ("O", 22),
+    ("X", 23),
+    ("Z", 23),
+    ("B", 23),
+    ("J", 23),
+    ("<cls>", 24),
+    ("<sep>", 25)])
+
+
+class TAPETokenizer():
+    r"""TAPE Tokenizer. Can use different vocabs depending on the model.
+    """
+
+    def __init__(self, vocab: str = 'iupac'):
+        if vocab == 'iupac':
+            self.vocab = IUPAC_VOCAB
+        elif vocab == 'unirep':
+            self.vocab = UNIREP_VOCAB
+        self.tokens = list(self.vocab.keys())
+        self._vocab_type = vocab
+        assert self.start_token in self.vocab and self.stop_token in self.vocab
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.vocab)
+
+    @property
+    def start_token(self) -> str:
+        return "<cls>"
+
+    @property
+    def stop_token(self) -> str:
+        return "<sep>"
+
+    @property
+    def mask_token(self) -> str:
+        if "<mask>" in self.vocab:
+            return "<mask>"
+        else:
+            raise RuntimeError(f"{self._vocab_type} vocab does not support masking")
+
+    def tokenize(self, text: str) -> List[str]:
+        return [x for x in text]
+
+    def convert_token_to_id(self, token: str) -> int:
+        """ Converts a token (str/unicode) in an id using the vocab. """
+        try:
+            return self.vocab[token]
+        except KeyError:
+            raise KeyError(f"Unrecognized token: '{token}'")
+
+    def convert_tokens_to_ids(self, tokens: List[str]) -> List[int]:
+        return [self.convert_token_to_id(token) for token in tokens]
+
+    def convert_id_to_token(self, index: int) -> str:
+        """Converts an index (integer) in a token (string/unicode) using the vocab."""
+        try:
+            return self.tokens[index]
+        except IndexError:
+            raise IndexError(f"Unrecognized index: '{index}'")
+
+    def convert_ids_to_tokens(self, indices: List[int]) -> List[str]:
+        return [self.convert_id_to_token(id_) for id_ in indices]
+
+    def convert_tokens_to_string(self, tokens: str) -> str:
+        """ Converts a sequence of tokens (string) in a single string. """
+        return ''.join(tokens)
+
+    def add_special_tokens(self, token_ids: List[str]) -> List[str]:
+        """
+        Adds special tokens to the a sequence for sequence classification tasks.
+        A BERT sequence has the following format: [CLS] X [SEP]
+        """
+        cls_token = [self.start_token]
+        sep_token = [self.stop_token]
+        return cls_token + token_ids + sep_token
+
+    def encode(self, text: str) -> np.ndarray:
+        tokens = self.tokenize(text)
+        tokens = self.add_special_tokens(tokens)
+        token_ids = self.convert_tokens_to_ids(tokens)
+        return np.array(token_ids, np.int64)
+
+    @classmethod
+    def from_pretrained(cls, **kwargs):
+        return cls()

tape/training.py

@@ -0,0 +1,659 @@
+import typing
+import os
+import logging
+from timeit import default_timer as timer
+import json
+from pathlib import Path
+import inspect
+import pickle as pkl
+
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from .optimization import WarmupLinearSchedule
+
+from . import utils
+from . import errors
+from . import visualization
+from .registry import registry
+from .models.modeling_utils import ProteinModel
+
+try:
+    from apex import amp
+    import amp_C
+    import apex_C
+    from apex.amp import _amp_state
+    from apex.parallel.distributed import flat_dist_call
+    from apex.parallel.distributed import DistributedDataParallel as DDP
+    APEX_FOUND = True
+except ImportError:
+    APEX_FOUND = False
+
+logger = logging.getLogger(__name__)
+
+MetricsDict = typing.Dict[str, float]
+LossAndMetrics = typing.Tuple[float, MetricsDict]
+OutputDict = typing.Dict[str, typing.Any]
+
+
+class ForwardRunner:
+
+    def __init__(self,
+                 model: ProteinModel,
+                 device: torch.device = torch.device('cuda:0'),
+                 n_gpu: int = 1,
+                 fp16: bool = False,
+                 local_rank: int = -1):
+
+        self.model = model
+        self.device = device
+        self.n_gpu = n_gpu
+        self.fp16 = fp16
+        self.local_rank = local_rank
+
+        forward_arg_keys = inspect.getfullargspec(model.forward).args
+        forward_arg_keys = forward_arg_keys[1:]  # remove self argument
+        self._forward_arg_keys = forward_arg_keys
+        assert 'input_ids' in self._forward_arg_keys
+
+    def initialize_distributed_model(self):
+        if self.local_rank != -1:
+            if not self.fp16:
+                self.model = DDP(self.model)
+            else:
+                flat_dist_call([param.data for param in self.model.parameters()],
+                               torch.distributed.broadcast, (0,))
+        elif self.n_gpu > 1:
+            self.model = nn.DataParallel(self.model)
+
+    def forward(self,
+                batch: typing.Dict[str, torch.Tensor],
+                return_outputs: bool = False,
+                no_loss: bool = False):
+        # Filter out batch items that aren't used in this model
+        # Requires that dataset keys match the forward args of the model
+        # Useful if some elements of the data are only used by certain models
+        # e.g. PSSMs / MSAs and other evolutionary data
+        batch = {name: tensor for name, tensor in batch.items()
+                 if name in self._forward_arg_keys}
+        if self.device.type == 'cuda':
+            batch = {name: tensor.cuda(device=self.device, non_blocking=True)
+                     for name, tensor in batch.items()}
+
+        outputs = self.model(**batch)
+
+        if no_loss:
+            return outputs
+
+        if isinstance(outputs[0], tuple):
+            # model also returned metrics
+            loss, metrics = outputs[0]
+        else:
+            # no metrics
+            loss = outputs[0]
+            metrics = {}
+
+        if self.n_gpu > 1:  # pytorch DataDistributed doesn't mean scalars
+            loss = loss.mean()
+            metrics = {name: metric.mean() for name, metric in metrics.items()}
+
+        if return_outputs:
+            return loss, metrics, outputs
+        else:
+            return loss, metrics
+
+    def train(self):
+        self.model.train()
+        return self
+
+    def eval(self):
+        self.model.eval()
+        return self
+
+
+class BackwardRunner(ForwardRunner):
+
+    def __init__(self,
+                 model: ProteinModel,
+                 optimizer: optim.Optimizer,  # type: ignore
+                 gradient_accumulation_steps: int = 1,
+                 device: torch.device = torch.device('cuda:0'),
+                 n_gpu: int = 1,
+                 fp16: bool = False,
+                 local_rank: int = -1,
+                 max_grad_norm: float = 1.0,
+                 warmup_steps: int = 0,
+                 num_train_optimization_steps: int = 1000000):
+
+        super().__init__(model, device, n_gpu, fp16, local_rank)
+        self.optimizer = optimizer
+        self.max_grad_norm = max_grad_norm
+        self._global_step = 0
+        self._local_rank = local_rank
+        self._overflow_buf = torch.cuda.IntTensor([0])  # type: ignore
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self._delay_accumulation = fp16 and local_rank != -1
+
+        self.scheduler = WarmupLinearSchedule(
+            self.optimizer, warmup_steps, num_train_optimization_steps)
+
+    def initialize_fp16(self):
+        if self.fp16:
+            self.model, self.optimizer = amp.initialize(
+                self.model, self.optimizer, opt_level="O2", loss_scale="dynamic",
+                master_weights=True)
+            _amp_state.loss_scalers[0]._loss_scale = 2 ** 20
+
+    def resume_from_checkpoint(self, checkpoint_dir: str) -> int:
+        checkpoint = torch.load(
+            os.path.join(checkpoint_dir, 'checkpoint.bin'), map_location=self.device)
+        self.optimizer.load_state_dict(checkpoint['optimizer'])
+        if self.fp16:
+            self.optimizer._lazy_init_maybe_master_weights()
+            self.optimizer._amp_stash.lazy_init_called = True
+            self.optimizer.load_state_dict(checkpoint['optimizer'])
+            for param, saved in zip(
+                    amp.master_params(self.optimizer), checkpoint['master params']):
+                param.data.copy_(saved.data)
+            amp.load_state_dict(checkpoint['amp'])
+        self.scheduler.load_state_dict(checkpoint['scheduler'])
+        start_epoch = checkpoint['epoch'] + 1
+        return start_epoch
+
+    def save_state(self, save_directory: typing.Union[str, Path], epoch_id: int):
+        save_directory = Path(save_directory)
+        if not save_directory.exists():
+            save_directory.mkdir()
+        else:
+            assert save_directory.is_dir(), "Save path should be a directory"
+        model_to_save = getattr(self.model, 'module', self.model)
+        model_to_save.save_pretrained(save_directory)
+        optimizer_state: typing.Dict[str, typing.Any] = {
+            'optimizer': self.optimizer.state_dict(),
+            'scheduler': self.scheduler.state_dict(),
+            'epoch': epoch_id}
+        if APEX_FOUND:
+            optimizer_state['master params'] = list(amp.master_params(self.optimizer))
+            try:
+                optimizer_state['amp'] = amp.state_dict()
+            except AttributeError:
+                pass
+        torch.save(optimizer_state, save_directory / 'checkpoint.bin')
+
+    def backward(self, loss) -> None:
+        if not self._delay_accumulation:
+            loss = loss / self.gradient_accumulation_steps
+        if self.fp16:
+            with amp.scale_loss(loss, self.optimizer,
+                                delay_overflow_check=self._delay_accumulation) as scaled_loss:
+                scaled_loss.backward()
+        else:
+            loss.backward()
+
+    def step(self) -> None:
+        nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+        if self._local_rank == -1:
+            self._step()
+        elif not self.fp16:
+            # TODO: Can you do this allreduce after accumulation also?
+            self._step()
+        else:
+            self._step_distributed_fp16()
+
+    def _step(self) -> None:
+        self.optimizer.step()
+        if self.scheduler is not None:
+            self.scheduler.step()  # type: ignore
+        self._global_step += 1
+
+    def _step_distributed_fp16(self) -> None:
+        # manually allreduce gradients after all accumulation steps
+        # check for Inf/NaN
+        # 1. allocate an uninitialized buffer for flattened gradient
+        scaler = _amp_state.loss_scalers[0]
+        master_grads = [p.grad for p in amp.master_params(self.optimizer) if p.grad is not None]
+        flat_grad_size = sum(p.numel() for p in master_grads)
+        # allreduce_dtype = torch.float16 if args.allreduce_post_accumulation_fp16 else \
+            # torch.float32
+        allreduce_dtype = torch.float16
+        flat_raw = torch.empty(flat_grad_size, device='cuda', dtype=allreduce_dtype)
+        # 2. combine unflattening and predivision of unscaled 'raw' gradient
+        allreduced_views = apex_C.unflatten(flat_raw, master_grads)
+        self._overflow_buf.zero_()
+        amp_C.multi_tensor_scale(
+            65536,
+            self._overflow_buf,
+            [master_grads, allreduced_views],
+            scaler.loss_scale() / (
+                torch.distributed.get_world_size() * self.gradient_accumulation_steps))
+        # 3. sum gradient across ranks. Because of the predivision, this averages the gradient
+        torch.distributed.all_reduce(flat_raw)
+        # 4. combine unscaling and unflattening of allreduced gradient
+        self._overflow_buf.zero_()
+        amp_C.multi_tensor_scale(
+            65536,
+            self._overflow_buf,
+            [allreduced_views, master_grads],
+            1. / scaler.loss_scale())
+        # 5. update loss scale
+        scaler = _amp_state.loss_scalers[0]
+        old_overflow_buf = scaler._overflow_buf
+        scaler._overflow_buf = self._overflow_buf
+        had_overflow = scaler.update_scale()
+        scaler._overfloat_buf = old_overflow_buf
+        # 6. call optimizer step function
+        if had_overflow == 0:
+            self._step()
+        else:
+            # Overflow detected, print message and clear gradients
+            logger.info(f"Gradient overflow.  Skipping step, reducing loss scale to "
+                        f"{scaler.loss_scale()}")
+            if _amp_state.opt_properties.master_weights:
+                for param in self.optimizer._amp_stash.all_fp32_from_fp16_params:
+                    param.grad = None
+        for param in self.model.parameters():
+            param.grad = None
+
+    @property
+    def global_step(self) -> int:
+        return self._global_step
+
+
+def run_train_epoch(epoch_id: int,
+                    train_loader: DataLoader,
+                    runner: BackwardRunner,
+                    viz: typing.Optional[visualization.TAPEVisualizer] = None,
+                    num_log_iter: int = 20,
+                    gradient_accumulation_steps: int = 1,
+                    num_steps_per_epoch: int = -1) -> LossAndMetrics:
+    if viz is None:
+        viz = visualization.DummyVisualizer()
+    smoothing = 1 - 1 / num_log_iter
+    accumulator = utils.MetricsAccumulator(smoothing)
+
+    torch.set_grad_enabled(True)
+    runner.train()
+
+    def make_log_str(step: int, time: float) -> str:
+        ep_percent = epoch_id + step / len(train_loader)
+        if runner.scheduler is not None:
+            curr_lr = runner.scheduler.get_lr()[0]  # type: ignore
+        else:
+            curr_lr = runner.optimizer.param_groups[0]['lr']
+
+        print_str = []
+        print_str.append(f"[Ep: {ep_percent:.2f}]")
+        print_str.append(f"[Iter: {runner.global_step}]")
+        print_str.append(f"[Time: {time:5.2f}s]")
+        print_str.append(f"[Loss: {accumulator.loss():.5g}]")
+
+        for name, value in accumulator.metrics().items():
+            print_str.append(f"[{name.capitalize()}: {value:.5g}]")
+
+        print_str.append(f"[LR: {curr_lr:.5g}]")
+        return ''.join(print_str)
+
+    start_t = timer()
+    for step, batch in enumerate(train_loader):
+        loss, metrics = runner.forward(batch)  # type: ignore
+        runner.backward(loss)
+        accumulator.update(loss, metrics, step=False)
+        if (step + 1) % gradient_accumulation_steps == 0:
+            runner.step()
+            viz.log_metrics(accumulator.step(), "train", runner.global_step)
+            if runner.global_step % num_log_iter == 0:
+                end_t = timer()
+                logger.info(make_log_str(step, end_t - start_t))
+                start_t = end_t
+        if num_steps_per_epoch != -1 and (step + 1) > num_steps_per_epoch:
+            break
+
+    final_print_str = f"Train: [Loss: {accumulator.final_loss():.5g}]"
+    for name, value in accumulator.final_metrics().items():
+        final_print_str += f"[{name.capitalize()}: {value:.5g}]"
+    logger.info(final_print_str)
+    return accumulator.final_loss(), accumulator.final_metrics()
+
+
+def run_valid_epoch(epoch_id: int,
+                    valid_loader: DataLoader,
+                    runner: ForwardRunner,
+                    viz: typing.Optional[visualization.TAPEVisualizer] = None,
+                    is_master: bool = True,
+                    val_check_frac: float = 1.0) -> typing.Tuple[float, typing.Dict[str, float]]:
+    num_batches = len(valid_loader)
+    num_batches_to_run = int(num_batches * val_check_frac)
+    accumulator = utils.MetricsAccumulator()
+
+    torch.set_grad_enabled(False)
+    runner.eval()
+
+    for idx, batch in enumerate(tqdm(valid_loader, desc='Running Eval', total=num_batches_to_run,
+                                disable=not is_master, leave=False)):
+        loss, metrics = runner.forward(batch)  # type: ignore
+        accumulator.update(loss, metrics)
+        if idx>num_batches_to_run:
+            break
+
+    # Reduce loss across all processes if multiprocessing
+    eval_loss = utils.reduce_scalar(accumulator.final_loss())
+    metrics = {name: utils.reduce_scalar(value)
+               for name, value in accumulator.final_metrics().items()}
+
+    print_str = f"Evaluation: [Loss: {eval_loss:.5g}]"
+    for name, value in metrics.items():
+        print_str += f"[{name.capitalize()}: {value:.5g}]"
+
+    metrics['loss'] = eval_loss
+    if viz is not None:
+        viz.log_metrics(metrics, "val", getattr(runner, 'global_step', epoch_id))
+
+    logger.info(print_str)
+
+    return eval_loss, metrics
+
+
+def _get_outputs_to_save(batch, outputs):
+    targets = batch['targets'].cpu().numpy()
+    outputs = outputs.cpu().numpy()
+    protein_length = batch['protein_length'].sum(1).cpu().numpy()
+
+    reshaped_output = []
+    for target, output, plength in zip(targets, outputs, protein_length):
+        output_slices = tuple(slice(1, plength - 1) if dim == protein_length.max() else
+                              slice(0, dim) for dim in output.shape)
+        output = output[output_slices]
+        target = target[output_slices]
+
+        reshaped_output.append((target, output))
+    reshaped_output
+
+
+def run_eval_epoch(eval_loader: DataLoader,
+                   runner: ForwardRunner,
+                   is_master: bool = True) -> typing.List[typing.Dict[str, typing.Any]]:
+    torch.set_grad_enabled(False)
+    runner.eval()
+
+    save_outputs = []
+
+    for batch in tqdm(eval_loader, desc='Evaluation', total=len(eval_loader),
+                      disable=not is_master):
+        loss, metrics, outputs = runner.forward(batch, return_outputs=True)  # type: ignore
+        predictions = outputs[1].cpu().numpy()
+        targets = batch['targets'].cpu().numpy()
+        for pred, target in zip(predictions, targets):
+            save_outputs.append({'prediction': pred, 'target': target})
+
+    return save_outputs
+
+
+def run_train(model_type: str,
+              task: str,
+              learning_rate: float = 1e-4,
+              batch_size: int = 1024,
+              num_train_epochs: int = 10,
+              num_log_iter: int = 20,
+              fp16: bool = False,
+              warmup_steps: int = 10000,
+              gradient_accumulation_steps: int = 1,
+              loss_scale: int = 0,
+              max_grad_norm: float = 1.0,
+              exp_name: typing.Optional[str] = None,
+              from_pretrained: typing.Optional[str] = None,
+              log_dir: str = './logs',
+              eval_freq: int = 1,
+              save_freq: typing.Union[int, str] = 1,
+              model_config_file: typing.Optional[str] = None,
+              data_dir: str = './data',
+              output_dir: str = './results',
+              no_cuda: bool = False,
+              seed: int = 42,
+              local_rank: int = -1,
+              tokenizer: str = 'iupac',
+              num_workers: int = 8,
+              debug: bool = False,
+              log_level: typing.Union[str, int] = logging.INFO,
+              patience: int = -1,
+              resume_from_checkpoint: bool = False,
+              model_args = None,
+              num_steps_per_epoch: int = -1,
+              val_check_frac: float = 1.0) -> None:
+
+    # SETUP AND LOGGING CODE #
+    input_args = locals()
+    device, n_gpu, is_master = utils.setup_distributed(
+        local_rank, no_cuda)
+
+    exp_dir = utils.get_expname(exp_name, task, model_type)
+    save_path = Path(output_dir) / exp_dir
+
+    if is_master:
+        # save all the hidden parameters.
+        save_path.mkdir(parents=True, exist_ok=True)
+        with (save_path / 'args.json').open('w') as f:
+            json.dump(input_args, f)
+
+    utils.barrier_if_distributed()
+    utils.setup_logging(local_rank, save_path, log_level)
+    utils.set_random_seeds(seed, n_gpu)
+
+    train_dataset = utils.setup_dataset(task, data_dir, 'train', tokenizer)
+    valid_dataset = utils.setup_dataset(task, data_dir, 'valid', tokenizer)
+    train_loader = utils.setup_loader(
+        train_dataset, batch_size, local_rank, n_gpu,
+        gradient_accumulation_steps, num_workers)
+    valid_loader = utils.setup_loader(
+        valid_dataset, batch_size, local_rank, n_gpu,
+        gradient_accumulation_steps, num_workers)
+
+    num_train_optimization_steps = utils.get_num_train_optimization_steps(
+        train_dataset, batch_size, num_train_epochs)
+
+    model = registry.get_task_model(model_type, task, model_config_file, from_pretrained, model_args)
+    model = model.to(device)
+    optimizer = utils.setup_optimizer(model, learning_rate)
+    viz = visualization.get(log_dir, exp_dir, local_rank, debug=debug)
+    viz.log_config(input_args)
+    viz.log_config(model.config.to_dict())
+    viz.watch(model)
+
+    logger.info(
+        f"device: {device} "
+        f"n_gpu: {n_gpu}, "
+        f"distributed_training: {local_rank != -1}, "
+        f"16-bits training: {fp16}")
+
+    runner = BackwardRunner(
+        model, optimizer, gradient_accumulation_steps, device, n_gpu,
+        fp16, local_rank, max_grad_norm, warmup_steps, num_train_optimization_steps)
+
+    runner.initialize_fp16()
+    if resume_from_checkpoint:
+        assert from_pretrained is not None
+        start_epoch = runner.resume_from_checkpoint(from_pretrained)
+    else:
+        start_epoch = 0
+    runner.initialize_distributed_model()
+
+    num_train_optimization_steps = utils.get_num_train_optimization_steps(
+        train_dataset, batch_size, num_train_epochs)
+    is_master = local_rank in (-1, 0)
+
+    if isinstance(save_freq, str) and save_freq != 'improvement':
+        raise ValueError(
+            f"Only recongized string value for save_freq is 'improvement'"
+            f", received: {save_freq}")
+
+    if save_freq == 'improvement' and eval_freq <= 0:
+        raise ValueError("Cannot set save_freq to 'improvement' and eval_freq < 0")
+
+    num_trainable_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    logger.info("***** Running training *****")
+    logger.info("  Num examples = %d", len(train_dataset))
+    logger.info("  Batch size = %d", batch_size)
+    logger.info("  Num epochs = %d", num_train_epochs)
+    logger.info("  Num train steps = %d", num_train_optimization_steps)
+    logger.info("  Num parameters = %d", num_trainable_parameters)
+
+    best_val_loss = float('inf')
+    num_evals_no_improvement = 0
+
+    def do_save(epoch_id: int, num_evals_no_improvement: int) -> bool:
+        if not is_master:
+            return False
+        if isinstance(save_freq, int):
+            return ((epoch_id + 1) % save_freq == 0) or ((epoch_id + 1) == num_train_epochs)
+        else:
+            return num_evals_no_improvement == 0
+
+    utils.barrier_if_distributed()
+
+    # ACTUAL TRAIN/EVAL LOOP #
+    with utils.wrap_cuda_oom_error(local_rank, batch_size, n_gpu, gradient_accumulation_steps):
+        for epoch_id in range(start_epoch, num_train_epochs):
+            run_train_epoch(epoch_id, train_loader, runner,
+                            viz, num_log_iter, gradient_accumulation_steps, num_steps_per_epoch)
+            if eval_freq > 0 and (epoch_id + 1) % eval_freq == 0:
+                val_loss, _ = run_valid_epoch(epoch_id, valid_loader, runner, viz, is_master, val_check_frac)
+                if val_loss < best_val_loss:
+                    best_val_loss = val_loss
+                    num_evals_no_improvement = 0
+                else:
+                    num_evals_no_improvement += 1
+
+            # Save trained model
+            if do_save(epoch_id, num_evals_no_improvement):
+                logger.info("** ** * Saving trained model ** ** * ")
+                # Only save the model itself
+                runner.save_state(save_path, epoch_id)
+                logger.info(f"Saving model checkpoint to {save_path}")
+
+            utils.barrier_if_distributed()
+            if patience > 0 and num_evals_no_improvement >= patience:
+                logger.info(f"Finished training at epoch {epoch_id} because no "
+                            f"improvement for {num_evals_no_improvement} epochs.")
+                logger.log(35, f"Best Val Loss: {best_val_loss}")
+                if local_rank != -1:
+                    # If you're distributed, raise this error. It sends a signal to
+                    # the master process which lets it kill other processes and terminate
+                    # without actually reporting an error. See utils/distributed_utils.py
+                    # for the signal handling code.
+                    raise errors.EarlyStopping
+                else:
+                    break
+    logger.info(f"Finished training after {num_train_epochs} epochs.")
+    if best_val_loss != float('inf'):
+        logger.log(35, f"Best Val Loss: {best_val_loss}")
+
+
+def run_eval(model_type: str,
+             task: str,
+             from_pretrained: str,
+             split: str = 'test',
+             batch_size: int = 1024,
+             model_config_file: typing.Optional[str] = None,
+             data_dir: str = './data',
+             no_cuda: bool = False,
+             seed: int = 42,
+             tokenizer: str = 'iupac',
+             num_workers: int = 8,
+             debug: bool = False,
+             metrics: typing.Tuple[str, ...] = (),
+             log_level: typing.Union[str, int] = logging.INFO) -> typing.Dict[str, float]:
+
+    local_rank = -1  # TAPE does not support torch.distributed.launch for evaluation
+    device, n_gpu, is_master = utils.setup_distributed(local_rank, no_cuda)
+    utils.setup_logging(local_rank, save_path=None, log_level=log_level)
+    utils.set_random_seeds(seed, n_gpu)
+
+    pretrained_dir = Path(from_pretrained)
+
+    logger.info(
+        f"device: {device} "
+        f"n_gpu: {n_gpu}")
+
+    model = registry.get_task_model(model_type, task, model_config_file, from_pretrained)
+    model = model.to(device)
+
+    runner = ForwardRunner(model, device, n_gpu)
+    runner.initialize_distributed_model()
+    valid_dataset = utils.setup_dataset(task, data_dir, split, tokenizer)
+    valid_loader = utils.setup_loader(
+        valid_dataset, batch_size, local_rank, n_gpu,
+        1, num_workers)
+
+    metric_functions = [registry.get_metric(name) for name in metrics]
+    save_outputs = run_eval_epoch(valid_loader, runner, is_master)
+    target = [el['target'] for el in save_outputs]
+    prediction = [el['prediction'] for el in save_outputs]
+
+    metrics_to_save = {name: metric(target, prediction)
+                       for name, metric in zip(metrics, metric_functions)}
+    logger.info(''.join(f'{name}: {val}' for name, val in metrics_to_save.items()))
+
+    with (pretrained_dir / 'results.pkl').open('wb') as f:
+        pkl.dump((metrics_to_save, save_outputs), f)
+
+    return metrics_to_save
+
+
+def run_embed(model_type: str,
+              data_file: str,
+              out_file: str,
+              from_pretrained: str,
+              batch_size: int = 1024,
+              model_config_file: typing.Optional[str] = None,
+              full_sequence_embed: bool = False,
+              no_cuda: bool = False,
+              seed: int = 42,
+              tokenizer: str = 'iupac',
+              num_workers: int = 8,
+              log_level: typing.Union[str, int] = logging.INFO) -> None:
+
+    local_rank = -1  # TAPE does not support torch.distributed.launch for embedding
+    device, n_gpu, is_master = utils.setup_distributed(local_rank, no_cuda)
+    utils.setup_logging(local_rank, save_path=None, log_level=log_level)
+    utils.set_random_seeds(seed, n_gpu)
+
+    logger.info(
+        f"device: {device} "
+        f"n_gpu: {n_gpu}")
+
+    task_spec = registry.get_task_spec('embed')
+    model = registry.get_task_model(
+        model_type, task_spec.name, model_config_file, from_pretrained)
+    model = model.to(device)
+    runner = ForwardRunner(model, device, n_gpu)
+    runner.initialize_distributed_model()
+    runner.eval()
+    torch.set_grad_enabled(False)
+
+    dataset = task_spec.dataset(data_file, tokenizer=tokenizer)  # type: ignore
+    valid_loader = utils.setup_loader(dataset, batch_size, local_rank, n_gpu, 1, num_workers)
+
+    with utils.IncrementalNPZ(out_file) as npzfile:
+        with utils.wrap_cuda_oom_error(local_rank, batch_size, n_gpu):
+            for batch in tqdm(valid_loader, total=len(valid_loader)):
+                outputs = runner.forward(batch, no_loss=True)
+                ids = batch['ids']
+                sequence_embed = outputs[0]
+                pooled_embed = outputs[1]
+                sequence_lengths = batch['input_mask'].sum(1)
+                sequence_embed = sequence_embed.cpu().numpy()
+                pooled_embed = pooled_embed.cpu().numpy()
+                sequence_lengths = sequence_lengths.cpu().numpy()
+
+                for seqembed, poolembed, length, protein_id in zip(
+                        sequence_embed, pooled_embed, sequence_lengths, ids):
+                    seqembed = seqembed[:length]
+                    arrays = {'pooled': poolembed}
+                    if not full_sequence_embed:
+                        # avgpool across the sequence
+                        arrays['avg'] = seqembed.mean(0)
+                    else:
+                        arrays['seq'] = seqembed
+                    to_save = {protein_id: arrays}
+                    npzfile.savez(**to_save)

tape/utils/__init__.py

@@ -0,0 +1,23 @@
+from .utils import int_or_str  # noqa: F401
+from .utils import check_is_file  # noqa: F401
+from .utils import check_is_dir  # noqa: F401
+from .utils import path_to_datetime  # noqa: F401
+from .utils import get_expname  # noqa: F401
+from .utils import get_effective_num_gpus  # noqa: F401
+from .utils import get_effective_batch_size  # noqa: F401
+from .utils import get_num_train_optimization_steps  # noqa: F401
+from .utils import set_random_seeds  # noqa: F401
+from .utils import MetricsAccumulator  # noqa: F401
+from .utils import wrap_cuda_oom_error  # noqa: F401
+from .utils import write_lmdb  # noqa: F401
+from .utils import IncrementalNPZ  # noqa: F401
+
+from .setup_utils import setup_logging  # noqa: F401
+from .setup_utils import setup_optimizer  # noqa: F401
+from .setup_utils import setup_dataset  # noqa: F401
+from .setup_utils import setup_loader  # noqa: F401
+from .setup_utils import setup_distributed  # noqa: F401
+
+from .distributed_utils import barrier_if_distributed  # noqa: F401
+from .distributed_utils import reduce_scalar  # noqa: F401
+from .distributed_utils import launch_process_group  # noqa: F401

tape/utils/_sampler.py

@@ -0,0 +1,98 @@
+"""Implementation of a bucketed data sampler from PyTorch-NLP.
+Modified by Roshan Rao.
+
+See https://github.com/PetrochukM/PyTorch-NLP/
+"""
+import typing
+import math
+import operator
+from torch.utils.data.sampler import Sampler
+from torch.utils.data.sampler import BatchSampler
+from torch.utils.data.sampler import SubsetRandomSampler
+
+
+class SortedSampler(Sampler):
+    """ Samples elements sequentially, always in the same order.
+    Args:
+        data (iterable): Iterable data.
+        sort_key (callable): Specifies a function of one argument that is used to extract a
+            numerical comparison key from each list element.
+    Example:
+        >>> list(SortedSampler(range(10), sort_key=lambda i: -i))
+        [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
+    """
+
+    def __init__(self,
+                 dataset,
+                 sort_key: typing.Callable[[int], typing.Any],
+                 indices: typing.Optional[typing.Iterable[int]] = None):
+        super().__init__(dataset)
+        self.dataset = dataset
+        self.sort_key = sort_key
+        if indices is None:
+            sort_keys = map(sort_key, dataset)
+        else:
+            sort_keys = ((i, sort_key(dataset[i])) for i in indices)
+        self.sorted_indices = [i for i, _ in sorted(sort_keys, key=operator.itemgetter(1))]
+
+    def __iter__(self):
+        return iter(self.sorted_indices)
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+class BucketBatchSampler(BatchSampler):
+    """ `BucketBatchSampler` toggles between `sampler` batches and sorted batches.
+    Typically, the `sampler` will be a `RandomSampler` allowing the user to toggle between
+    random batches and sorted batches. A larger `bucket_size_multiplier` is more sorted
+    and vice versa. Provides ~10-25 percent speedup.
+
+    Background:
+        ``BucketBatchSampler`` is similar to a ``BucketIterator`` found in popular
+        libraries like ``AllenNLP`` and ``torchtext``. A ``BucketIterator`` pools together
+        examples with a similar size length to reduce the padding required for each batch
+        while maintaining some noise through bucketing.
+
+    Args:
+        sampler (torch.data.utils.sampler.Sampler):
+        batch_size (int): Size of mini-batch.
+        drop_last (bool): If `True` the sampler will drop the last batch if its size
+            would be less than `batch_size`.
+        sort_key (callable, optional): Callable to specify a comparison key for sorting.
+        bucket_size_multiplier (int, optional): Buckets are of size
+            `batch_size * bucket_size_multiplier`.
+    Example:
+        >>> from torch.utils.data.sampler import SequentialSampler
+        >>> sampler = SequentialSampler(list(range(10)))
+        >>> list(BucketBatchSampler(sampler, batch_size=3, drop_last=False))
+        [[6, 7, 8], [0, 1, 2], [3, 4, 5], [9]]
+        >>> list(BucketBatchSampler(sampler, batch_size=3, drop_last=True))
+        [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
+    """
+
+    def __init__(self,
+                 sampler,
+                 batch_size,
+                 drop_last,
+                 sort_key,
+                 dataset,
+                 bucket_size_multiplier=100):
+        super().__init__(sampler, batch_size, drop_last)
+        self.sort_key = sort_key
+        self.dataset = dataset
+        self.bucket_sampler = BatchSampler(
+            sampler, min(batch_size * bucket_size_multiplier, len(sampler)), False)
+
+    def __iter__(self):
+        for bucket in self.bucket_sampler:
+            sorted_sampler = SortedSampler(self.dataset, self.sort_key, indices=bucket)
+            for batch in SubsetRandomSampler(
+                    list(BatchSampler(sorted_sampler, self.batch_size, self.drop_last))):
+                yield batch
+
+    def __len__(self):
+        if self.drop_last:
+            return len(self.sampler) // self.batch_size
+        else:
+            return math.ceil(len(self.sampler) / self.batch_size)

tape/utils/distributed_utils.py

@@ -0,0 +1,170 @@
+import typing
+import argparse
+import os
+import multiprocessing as mp
+import sys
+import signal
+
+import torch
+import torch.distributed as dist
+from torch.multiprocessing import _prctl_pr_set_pdeathsig  # type: ignore
+
+from ..errors import EarlyStopping
+
+
+def reduce_scalar(scalar: float) -> float:
+    if dist.is_available() and dist.is_initialized():
+        float_tensor = torch.cuda.FloatTensor([scalar])  # type: ignore
+        dist.all_reduce(float_tensor)
+        float_tensor /= dist.get_world_size()
+        scalar = float_tensor.item()
+    return scalar
+
+
+def barrier_if_distributed() -> None:
+    """Raises a barrier if in a distributed context, otherwise does nothing."""
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+
+
+def _wrap(fn, kwargs, error_queue):
+    # prctl(2) is a Linux specific system call.
+    # On other systems the following function call has no effect.
+    # This is set to ensure that non-daemonic child processes can
+    # terminate if their parent terminates before they do.
+    _prctl_pr_set_pdeathsig(signal.SIGINT)
+
+    try:
+        fn(**kwargs)
+    except KeyboardInterrupt:
+        pass  # SIGINT; Killed by parent, do nothing
+    except EarlyStopping:
+        sys.exit(signal.SIGUSR1)  # tape early stop exception
+    except Exception:
+        # Propagate exception to parent process, keeping original traceback
+        import traceback
+        error_queue.put(traceback.format_exc())
+        sys.exit(1)
+
+
+class ProcessContext:
+    def __init__(self, processes, error_queues):
+        self.error_queues = error_queues
+        self.processes = processes
+        self.sentinels = {
+            process.sentinel: index
+            for index, process in enumerate(processes)
+        }
+
+    def pids(self):
+        return [int(process.pid) for process in self.processes]
+
+    def join(self, timeout=None):
+        r"""
+        Tries to join one or more processes in this process context.
+        If one of them exited with a non-zero exit status, this function
+        kills the remaining processes and raises an exception with the cause
+        of the first process exiting.
+
+        Returns ``True`` if all processes have been joined successfully,
+        ``False`` if there are more processes that need to be joined.
+
+        Arguments:
+            timeout (float): Wait this long before giving up on waiting.
+        """
+        # Ensure this function can be called even when we're done.
+        if len(self.sentinels) == 0:
+            return True
+
+        # Wait for any process to fail or all of them to succeed.
+        ready = mp.connection.wait(
+            self.sentinels.keys(),
+            timeout=timeout,
+        )
+        error_index = None
+        for sentinel in ready:
+            index = self.sentinels.pop(sentinel)
+            process = self.processes[index]
+            process.join()
+            if process.exitcode != 0:
+                error_index = index
+                break
+        # Return if there was no error.
+        if error_index is None:
+            # Return whether or not all processes have been joined.
+            return len(self.sentinels) == 0
+        # Assume failure. Terminate processes that are still alive.
+        for process in self.processes:
+            if process.is_alive():
+                process.terminate()
+            process.join()
+
+        # There won't be an error on the queue if the process crashed.
+        if self.error_queues[error_index].empty():
+            exitcode = self.processes[error_index].exitcode
+            if exitcode == signal.SIGUSR1:
+                return True
+            elif exitcode < 0:
+                name = signal.Signals(-exitcode).name
+                raise Exception(
+                    "process %d terminated with signal %s" %
+                    (error_index, name)
+                )
+            else:
+                raise Exception(
+                    "process %d terminated with exit code %d" %
+                    (error_index, exitcode)
+                )
+
+        original_trace = self.error_queues[error_index].get()
+        msg = "\n\n-- Process %d terminated with the following error:\n" % error_index
+        msg += original_trace
+        raise Exception(msg)
+
+
+def launch_process_group(func: typing.Callable,
+                         args: argparse.Namespace,
+                         num_processes: int,
+                         num_nodes: int = 1,
+                         node_rank: int = 0,
+                         master_addr: str = "127.0.0.1",
+                         master_port: int = 29500,
+                         join: bool = True,
+                         daemon: bool = False):
+    # world size in terms of number of processes
+    dist_world_size = num_processes * num_nodes
+
+    # set PyTorch distributed related environmental variables
+    current_env = os.environ.copy()
+    current_env["MASTER_ADDR"] = master_addr
+    current_env["MASTER_PORT"] = str(master_port)
+    current_env["WORLD_SIZE"] = str(dist_world_size)
+    if 'OMP_NUM_THREADS' not in os.environ and num_processes > 1:
+        current_env["OMP_NUM_THREADS"] = str(4)
+
+    error_queues = []
+    processes = []
+
+    for local_rank in range(num_processes):
+        # each process's rank
+        dist_rank = num_processes * node_rank + local_rank
+        current_env["RANK"] = str(dist_rank)
+        current_env["LOCAL_RANK"] = str(local_rank)
+        args.local_rank = local_rank
+
+        error_queue: mp.SimpleQueue[Exception] = mp.SimpleQueue()
+        kwargs = {'args': args, 'env': current_env}
+        process = mp.Process(
+            target=_wrap,
+            args=(func, kwargs, error_queue),
+            daemon=daemon)
+        process.start()
+        error_queues.append(error_queue)
+        processes.append(process)
+
+    process_context = ProcessContext(processes, error_queues)
+    if not join:
+        return process_context
+
+    while not process_context.join():
+        pass

tape/utils/setup_utils.py

@@ -0,0 +1,135 @@
+"""Utility functions to help setup the model, optimizer, distributed compute, etc.
+"""
+import typing
+import logging
+from pathlib import Path
+import sys
+
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader, RandomSampler, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from ..optimization import AdamW
+
+from ..registry import registry
+
+from .utils import get_effective_batch_size
+from ._sampler import BucketBatchSampler
+
+logger = logging.getLogger(__name__)
+
+
+def setup_logging(local_rank: int,
+                  save_path: typing.Optional[Path] = None,
+                  log_level: typing.Union[str, int] = None) -> None:
+    if log_level is None:
+        level = logging.INFO
+    elif isinstance(log_level, str):
+        level = getattr(logging, log_level.upper())
+    elif isinstance(log_level, int):
+        level = log_level
+
+    if local_rank not in (-1, 0):
+        level = max(level, logging.WARN)
+
+    root_logger = logging.getLogger()
+    root_logger.setLevel(level)
+
+    formatter = logging.Formatter(
+        "%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%y/%m/%d %H:%M:%S")
+
+    if not root_logger.hasHandlers():
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setLevel(level)
+        console_handler.setFormatter(formatter)
+        root_logger.addHandler(console_handler)
+
+        if save_path is not None:
+            file_handler = logging.FileHandler(save_path / 'log')
+            file_handler.setLevel(level)
+            file_handler.setFormatter(formatter)
+            root_logger.addHandler(file_handler)
+
+
+def setup_optimizer(model,
+                    learning_rate: float):
+    """Create the AdamW optimizer for the given model with the specified learning rate. Based on
+    creation in the pytorch_transformers repository.
+
+    Args:
+        model (PreTrainedModel): The model for which to create an optimizer
+        learning_rate (float): Default learning rate to use when creating the optimizer
+
+    Returns:
+        optimizer (AdamW): An AdamW optimizer
+
+    """
+    param_optimizer = list(model.named_parameters())
+    no_decay = ['bias', 'gamma', 'beta', 'LayerNorm']
+    optimizer_grouped_parameters = [
+        {
+            "params": [
+                p for n, p in param_optimizer if not any(nd in n for nd in no_decay)
+            ],
+            "weight_decay": 0.01,
+        },
+        {
+            "params": [
+                p for n, p in param_optimizer if any(nd in n for nd in no_decay)
+            ],
+            "weight_decay": 0.0,
+        },
+    ]
+
+    optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate)
+    return optimizer
+
+
+def setup_dataset(task: str,
+                  data_dir: typing.Union[str, Path],
+                  split: str,
+                  tokenizer: str) -> Dataset:
+    task_spec = registry.get_task_spec(task)
+    return task_spec.dataset(data_dir, split, tokenizer)  # type: ignore
+
+
+def setup_loader(dataset: Dataset,
+                 batch_size: int,
+                 local_rank: int,
+                 n_gpu: int,
+                 gradient_accumulation_steps: int,
+                 num_workers: int) -> DataLoader:
+    sampler = DistributedSampler(dataset) if local_rank != -1 else RandomSampler(dataset)
+    batch_size = get_effective_batch_size(
+        batch_size, local_rank, n_gpu, gradient_accumulation_steps) * n_gpu
+    # WARNING: this will fail if the primary sequence is not the first thing the dataset returns
+    batch_sampler = BucketBatchSampler(
+        sampler, batch_size, False, lambda x: len(x[0]), dataset)
+
+    loader = DataLoader(
+        dataset,
+        num_workers=num_workers,
+        collate_fn=dataset.collate_fn,  # type: ignore
+        batch_sampler=batch_sampler)
+
+    return loader
+
+
+def setup_distributed(local_rank: int,
+                      no_cuda: bool) -> typing.Tuple[torch.device, int, bool]:
+    if local_rank != -1 and not no_cuda:
+        torch.cuda.set_device(local_rank)
+        device: torch.device = torch.device("cuda", local_rank)
+        n_gpu = 1
+        dist.init_process_group(backend="nccl")
+    elif not torch.cuda.is_available() or no_cuda:
+        device = torch.device("cpu")
+        n_gpu = 1
+    else:
+        device = torch.device("cuda")
+        n_gpu = torch.cuda.device_count()
+
+    is_master = local_rank in (-1, 0)
+
+    return device, n_gpu, is_master

tape/utils/utils.py

@@ -0,0 +1,327 @@
+import typing
+import random
+from pathlib import Path
+import logging
+from time import strftime, gmtime
+from datetime import datetime
+import os
+import argparse
+import contextlib
+from collections import defaultdict
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import torch.distributed as dist
+
+logger = logging.getLogger(__name__)
+FloatOrTensor = typing.Union[float, torch.Tensor]
+
+
+def int_or_str(arg: str) -> typing.Union[int, str]:
+    try:
+        return int(arg)
+    except ValueError:
+        return arg
+
+
+def check_is_file(file_path: str) -> str:
+    if file_path is None or os.path.isfile(file_path):
+        return file_path
+    else:
+        raise argparse.ArgumentTypeError(f"File path: {file_path} is not a valid file")
+
+
+def check_is_dir(dir_path: str) -> str:
+    if dir_path is None or os.path.isdir(dir_path):
+        return dir_path
+    else:
+        raise argparse.ArgumentTypeError(f"Directory path: {dir_path} is not a valid directory")
+
+
+def path_to_datetime(path: Path) -> datetime:
+    name = path.name
+    datetime_string = name.split('_')[0]
+    try:
+        year, month, day, hour, minute, second = datetime_string.split('-')
+    except ValueError:
+        try:
+            # Deprecated datetime strings
+            year, month, day, time_str = datetime_string.split('-')
+            hour, minute, second = time_str.split(':')
+        except ValueError:
+            return datetime(1, 1, 1)
+
+    pathdatetime = datetime(
+        int(year), int(month), int(day), int(hour), int(minute), int(second))
+    return pathdatetime
+
+
+def get_expname(exp_name: typing.Optional[str],
+                task: typing.Optional[str] = None,
+                model_type: typing.Optional[str] = None) -> str:
+    if exp_name is None:
+        time_stamp = strftime("%y-%m-%d-%H-%M-%S", gmtime())
+        exp_name = f"{task}_{model_type}_{time_stamp}_{random.randint(0, int(1e6)):0>6d}"
+    return exp_name
+
+
+def set_random_seeds(seed: int, n_gpu: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if n_gpu > 0:
+        torch.cuda.manual_seed_all(seed)  # type: ignore
+
+
+def get_effective_num_gpus(local_rank: int, n_gpu: int) -> int:
+    if local_rank == -1:
+        num_gpus = n_gpu
+    else:
+        num_gpus = dist.get_world_size()
+    return num_gpus
+
+
+def get_effective_batch_size(batch_size: int,
+                             local_rank: int,
+                             n_gpu: int,
+                             gradient_accumulation_steps: int = 1) -> int:
+    eff_batch_size = float(batch_size)
+    eff_batch_size /= gradient_accumulation_steps
+    eff_batch_size /= get_effective_num_gpus(local_rank, n_gpu)
+    return int(eff_batch_size)
+
+
+def get_num_train_optimization_steps(dataset: Dataset,
+                                     batch_size: int,
+                                     num_train_epochs: int) -> int:
+    return int(len(dataset) / batch_size * num_train_epochs)
+
+
+class MetricsAccumulator:
+
+    def __init__(self, smoothing: float = 0.95):
+        self._loss_tmp = 0.
+        self._smoothloss: typing.Optional[float] = None
+        self._totalloss = 0.
+        self._metricstmp: typing.Dict[str, float] = defaultdict(lambda: 0.0)
+        self._smoothmetrics: typing.Dict[str, float] = {}
+        self._totalmetrics: typing.Dict[str, float] = defaultdict(lambda: 0.0)
+
+        self._nacc_steps = 0
+        self._nupdates = 0
+        self._smoothing = smoothing
+
+    def update(self,
+               loss: FloatOrTensor,
+               metrics: typing.Dict[str, FloatOrTensor],
+               step: bool = True) -> None:
+        if isinstance(loss, torch.Tensor):
+            loss = loss.item()
+
+        self._loss_tmp += loss
+        for name, value in metrics.items():
+            if isinstance(value, torch.Tensor):
+                value = value.item()
+            self._metricstmp[name] += value
+        self._nacc_steps += 1
+
+        if step:
+            self.step()
+
+    def step(self) -> typing.Dict[str, float]:
+        loss_tmp = self._loss_tmp / self._nacc_steps
+        metricstmp = {name: value / self._nacc_steps
+                      for name, value in self._metricstmp.items()}
+
+        if self._smoothloss is None:
+            self._smoothloss = loss_tmp
+        else:
+            self._smoothloss *= self._smoothing
+            self._smoothloss += (1 - self._smoothing) * loss_tmp
+        self._totalloss += loss_tmp
+
+        for name, value in metricstmp.items():
+            if name in self._smoothmetrics:
+                currvalue = self._smoothmetrics[name]
+                newvalue = currvalue * self._smoothing + value * (1 - self._smoothing)
+            else:
+                newvalue = value
+
+            self._smoothmetrics[name] = newvalue
+            self._totalmetrics[name] += value
+
+        self._nupdates += 1
+
+        self._nacc_steps = 0
+        self._loss_tmp = 0
+        self._metricstmp = defaultdict(lambda: 0.0)
+
+        metricstmp['loss'] = loss_tmp
+        return metricstmp
+
+    def loss(self) -> float:
+        if self._smoothloss is None:
+            raise RuntimeError("Trying to get the loss without any updates")
+        return self._smoothloss
+
+    def metrics(self) -> typing.Dict[str, float]:
+        if self._nupdates == 0:
+            raise RuntimeError("Trying to get metrics without any updates")
+        return dict(self._smoothmetrics)
+
+    def final_loss(self) -> float:
+        return self._totalloss / self._nupdates
+
+    def final_metrics(self) -> typing.Dict[str, float]:
+        return {name: value / self._nupdates
+                for name, value in self._totalmetrics.items()}
+
+
+class wrap_cuda_oom_error(contextlib.ContextDecorator):
+    """A context manager that wraps the Cuda OOM message so that you get some more helpful
+    context as to what you can/should change. Can also be used as a decorator.
+
+    Examples:
+        1) As a context manager:
+
+            with wrap_cuda_oom_error(local_rank, batch_size, n_gpu, gradient_accumulation):
+                loss = model.forward(batch)
+                loss.backward()
+                optimizer.step()
+                optimizer.zero_grad
+
+        2) As a decorator:
+
+            @wrap_cuda_oom_error(local_rank, batch_size, n_gpu, gradient_accumulation)
+            def run_train_epoch(args):
+                ...
+                <code to run training epoch>
+                ...
+    """
+
+    def __init__(self,
+                 local_rank: int,
+                 batch_size: int,
+                 n_gpu: int = 1,
+                 gradient_accumulation_steps: typing.Optional[int] = None):
+        self._local_rank = local_rank
+        self._batch_size = batch_size
+        self._n_gpu = n_gpu
+        self._gradient_accumulation_steps = gradient_accumulation_steps
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        exc_args = exc_value.args if exc_value is not None else None
+        if exc_args and 'CUDA out of memory' in exc_args[0]:
+            eff_ngpu = get_effective_num_gpus(self._local_rank, self._n_gpu)
+            if self._gradient_accumulation_steps is not None:
+                eff_batch_size = get_effective_batch_size(
+                    self._batch_size, self._local_rank, self._n_gpu,
+                    self._gradient_accumulation_steps)
+                message = (f"CUDA out of memory. Reduce batch size or increase "
+                           f"gradient_accumulation_steps to divide each batch over more "
+                           f"forward passes.\n\n"
+                           f"\tHyperparameters:\n"
+                           f"\t\tbatch_size per backward-pass: {self._batch_size}\n"
+                           f"\t\tgradient_accumulation_steps: "
+                           f"{self._gradient_accumulation_steps}\n"
+                           f"\t\tn_gpu: {eff_ngpu}\n"
+                           f"\t\tbatch_size per (gpu * forward-pass): "
+                           f"{eff_batch_size}")
+            else:
+                eff_batch_size = get_effective_batch_size(
+                    self._batch_size, self._local_rank, self._n_gpu)
+                message = (f"CUDA out of memory. Reduce batch size to fit each "
+                           f"iteration in memory.\n\n"
+                           f"\tHyperparameters:\n"
+                           f"\t\tbatch_size per forward-pass: {self._batch_size}\n"
+                           f"\t\tn_gpu: {eff_ngpu}\n"
+                           f"\t\tbatch_size per (gpu * forward-pass): "
+                           f"{eff_batch_size}")
+            raise RuntimeError(message)
+        return False
+
+
+def write_lmdb(filename: str, iterable: typing.Iterable, map_size: int = 2 ** 20):
+    """Utility for writing a dataset to an LMDB file.
+
+    Args:
+        filename (str): Output filename to write to
+        iterable (Iterable): An iterable dataset to write to. Entries must be pickleable.
+        map_size (int, optional): Maximum allowable size of database in bytes. Required by LMDB.
+            You will likely have to increase this. Default: 1MB.
+    """
+    import lmdb
+    import pickle as pkl
+    env = lmdb.open(filename, map_size=map_size)
+
+    with env.begin(write=True) as txn:
+        for i, entry in enumerate(iterable):
+            txn.put(str(i).encode(), pkl.dumps(entry))
+        txn.put(b'num_examples', pkl.dumps(i + 1))
+    env.close()
+
+
+class IncrementalNPZ(object):
+    # Modified npz that allows incremental saving, from https://stackoverflow.com/questions/22712292/how-to-use-numpy-savez-in-a-loop-for-save-more-than-one-array  # noqa: E501
+    def __init__(self, file):
+        import tempfile
+        import zipfile
+        import os
+
+        if isinstance(file, str):
+            if not file.endswith('.npz'):
+                file = file + '.npz'
+
+        compression = zipfile.ZIP_STORED
+
+        zipfile = self.zipfile_factory(file, mode="a", compression=compression)
+
+        # Stage arrays in a temporary file on disk, before writing to zip.
+        fd, tmpfile = tempfile.mkstemp(suffix='-numpy.npy')
+        os.close(fd)
+
+        self.tmpfile = tmpfile
+        self.zip = zipfile
+        self._i = 0
+
+    def zipfile_factory(self, *args, **kwargs):
+        import zipfile
+        import sys
+        if sys.version_info >= (2, 5):
+            kwargs['allowZip64'] = True
+        return zipfile.ZipFile(*args, **kwargs)
+
+    def savez(self, *args, **kwds):
+        import os
+        import numpy.lib.format as fmt
+
+        namedict = kwds
+        for val in args:
+            key = 'arr_%d' % self._i
+            if key in namedict.keys():
+                raise ValueError("Cannot use un-named variables and keyword %s" % key)
+            namedict[key] = val
+            self._i += 1
+
+        try:
+            for key, val in namedict.items():
+                fname = key + '.npy'
+                fid = open(self.tmpfile, 'wb')
+                with open(self.tmpfile, 'wb') as fid:
+                    fmt.write_array(fid, np.asanyarray(val), allow_pickle=True)
+                self.zip.write(self.tmpfile, arcname=fname)
+        finally:
+            os.remove(self.tmpfile)
+
+    def close(self):
+        self.zip.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()

tape/visualization.py

@@ -0,0 +1,124 @@
+import typing
+import os
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+import torch.nn as nn
+
+from tensorboardX import SummaryWriter
+
+try:
+    import wandb
+    WANDB_FOUND = True
+except ImportError:
+    WANDB_FOUND = False
+
+logger = logging.getLogger(__name__)
+
+
+class TAPEVisualizer(ABC):
+    """Base class for visualization in TAPE"""
+
+    @abstractmethod
+    def __init__(self, log_dir: typing.Union[str, Path], exp_name: str, debug: bool = False):
+        raise NotImplementedError
+
+    @abstractmethod
+    def log_config(self, config: typing.Dict[str, typing.Any]) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def watch(self, model: nn.Module) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def log_metrics(self,
+                    metrics_dict: typing.Dict[str, float],
+                    split: str,
+                    step: int):
+        raise NotImplementedError
+
+
+class DummyVisualizer(TAPEVisualizer):
+    """Dummy class that doesn't do anything. Used for non-master branches."""
+
+    def __init__(self,
+                 log_dir: typing.Union[str, Path] = '',
+                 exp_name: str = '',
+                 debug: bool = False):
+        pass
+
+    def log_config(self, config: typing.Dict[str, typing.Any]) -> None:
+        pass
+
+    def watch(self, model: nn.Module) -> None:
+        pass
+
+    def log_metrics(self,
+                    metrics_dict: typing.Dict[str, float],
+                    split: str,
+                    step: int):
+        pass
+
+
+class TBVisualizer(TAPEVisualizer):
+
+    def __init__(self, log_dir: typing.Union[str, Path], exp_name: str, debug: bool = False):
+        log_dir = Path(log_dir) / exp_name
+        logger.info(f"tensorboard file at: {log_dir}")
+        self.logger = SummaryWriter(log_dir=str(log_dir))
+
+    def log_config(self, config: typing.Dict[str, typing.Any]) -> None:
+        logger.warn("Cannot log config when using a TBVisualizer. "
+                    "Configure wandb for this functionality")
+
+    def watch(self, model: nn.Module) -> None:
+        logger.warn("Cannot watch models when using a TBVisualizer. "
+                    "Configure wandb for this functionality")
+
+    def log_metrics(self,
+                    metrics_dict: typing.Dict[str, float],
+                    split: str,
+                    step: int):
+        for name, value in metrics_dict.items():
+            self.logger.add_scalar(split + "/" + name, value, step)
+
+
+class WandBVisualizer(TAPEVisualizer):
+
+    def __init__(self, log_dir: typing.Union[str, Path], exp_name: str, debug: bool = False):
+        if not WANDB_FOUND:
+            raise ImportError("wandb module not available")
+        #if debug:
+        #    os.environ['WANDB_MODE'] = 'dryrun'
+        #if 'WANDB_PROJECT' not in os.environ:
+        #    # Want the user to set the WANDB_PROJECT.
+        #    logger.warning("WANDB_PROJECT environment variable not found, "
+        #                   "not logging to app.wandb.ai")
+        #    os.environ['WANDB_MODE'] = 'dryrun'
+        wandb.init(dir=log_dir, name=exp_name)
+
+    def log_config(self, config: typing.Dict[str, typing.Any]) -> None:
+        wandb.config.update(config)
+
+    def watch(self, model: nn.Module):
+        wandb.watch(model)
+
+    def log_metrics(self,
+                    metrics_dict: typing.Dict[str, float],
+                    split: str,
+                    step: int):
+        wandb.log({f"{split.capitalize()} {name.capitalize()}": value
+                   for name, value in metrics_dict.items()}, step=step)
+
+
+def get(log_dir: typing.Union[str, Path],
+        exp_name: str,
+        local_rank: int,
+        debug: bool = False) -> TAPEVisualizer:
+    if local_rank not in (-1, 0):
+        return DummyVisualizer(log_dir, exp_name, debug)
+    elif WANDB_FOUND:
+        return WandBVisualizer(log_dir, exp_name, debug)
+    else:
+        return TBVisualizer(log_dir, exp_name, debug)