Add MIF and MIF-ST and increment version.

This commit is contained in:
Kevin Kaichuang Yang 2022-05-26 10:41:21 -04:00
Родитель f2be26244d
Коммит 7f7a1f68de
4 изменённых файлов: 558 добавлений и 25 удалений

Просмотреть файл

@ -8,9 +8,9 @@ Here we will demonstrate the application of several tools we hope will help with
pip install sequence-models
```
### Convolutional autoencoding representations of proteins (CARP)
### Loading pretrained models
We make available pretrained CNN protein sequence masked language models of various sizes. All of these have a ByteNet encoder architecture and are pretrained on the March 2020 release of UniRef50 using the same masked language modeling task as in BERT and ESM-1b. Models require PyTorch. We tested on `v1.9.0` and `v1.11.0`. If you installed into a clean conda environment, you may also need to install pandas, scipy, and wget.
Models require PyTorch. We tested on `v1.9.0` and `v1.11.0`. If you installed into a clean conda environment, you may also need to install pandas, scipy, and wget.
To load a model:
@ -20,19 +20,57 @@ from sequence_models.pretrained import load_model_and_alphabet
model, collater = load_model_and_alphabet('carp_640M')
```
The available models are `carp_600k`, `carp_38M`, `carp_76M`, and `carp_640M`.
Available models are
- `carp_600k`
- `carp_38M`
- `carp_76M`
- `carp_640M`
- `mif`
- `mifst`
You can also download the weights and hyperparameters manually from [Zenodo](https://doi.org/10.5281/zenodo.6368483).
### Convolutional autoencoding representations of proteins (CARP)
We make available pretrained CNN protein sequence masked language models of various sizes. All of these have a ByteNet encoder architecture and are pretrained on the March 2020 release of UniRef50 using the same masked language modeling task as in BERT and ESM-1b.
CARP is described in this [preprint](https://doi.org/10.1101/2022.05.19.492714).
You can also download the weights manually from [Zenodo](https://doi.org/10.5281/zenodo.6368483).
To encode a batch of sequences:
```
seqs = [['MDREQ'], ['MGTRRLLP']]
x = collater(seqs) # (n, max_len)
x = collater(seqs)[0] # (n, max_len)
rep = model(x) # (n, max_len, d_model)
```
The collater will pad sequences to the maximum length, and the model automatically ignores the padding.
### Masked Inverse Folding (MIF) and Masked Inverse Folding with Sequence Transfer (MIF-ST)
We make available pretrained masked inverse folding models with and without sequence pretraining transfer from CARP-640M.
[comment]: <> (MIF and MIF-ST are described in this [preprint]&#40;&#41;.)
You can also download the weights manually from [Zenodo](https://doi.org/10.1234/mifst).
To encode a sequence with its structure:
```
from sequence_models.pdb_utils import parse_PDB, process_coords
coords, wt, _ = parse_PDB('examples/gb1_a60fb_unrelaxed_rank_1_model_5.pdb')
coords = {
'N': coords[:, 0],
'CA': coords[:, 1],
'C': coords[:, 2]
}
dist, omega, theta, phi = process_coords(coords)
batch = [[wt, torch.tensor(dist, dtype=torch.float),
torch.tensor(omega, dtype=torch.float),
torch.tensor(theta, dtype=torch.float), torch.tensor(phi, dtype=torch.float)]]
src, nodes, edges, connections, edge_mask = collater(batch)
rep = model(src, nodes, edges, connections, edge_mask)
```
### Sequence Datasets and Dataloaders

Просмотреть файл

@ -0,0 +1,441 @@
MODEL 1
ATOM 1 N MET A 1 8.118 7.168 -10.526 1.00 93.37 N
ATOM 2 CA MET A 1 7.595 5.811 -10.659 1.00 93.37 C
ATOM 3 C MET A 1 6.146 5.738 -10.190 1.00 93.37 C
ATOM 4 CB MET A 1 8.453 4.824 -9.865 1.00 93.37 C
ATOM 5 O MET A 1 5.728 6.513 -9.327 1.00 93.37 O
ATOM 6 CG MET A 1 9.904 4.773 -10.315 1.00 93.37 C
ATOM 7 SD MET A 1 10.958 3.790 -9.179 1.00 93.37 S
ATOM 8 CE MET A 1 12.592 4.116 -9.896 1.00 93.37 C
ATOM 9 N GLN A 2 5.461 5.046 -10.984 1.00 97.03 N
ATOM 10 CA GLN A 2 4.054 4.922 -10.617 1.00 97.03 C
ATOM 11 C GLN A 2 3.792 3.613 -9.878 1.00 97.03 C
ATOM 12 CB GLN A 2 3.165 5.013 -11.858 1.00 97.03 C
ATOM 13 O GLN A 2 4.213 2.545 -10.328 1.00 97.03 O
ATOM 14 CG GLN A 2 1.677 5.082 -11.545 1.00 97.03 C
ATOM 15 CD GLN A 2 0.824 5.279 -12.784 1.00 97.03 C
ATOM 16 NE2 GLN A 2 -0.489 5.353 -12.595 1.00 97.03 N
ATOM 17 OE1 GLN A 2 1.342 5.363 -13.902 1.00 97.03 O
ATOM 18 N TYR A 3 3.102 3.739 -8.670 1.00 98.33 N
ATOM 19 CA TYR A 3 2.745 2.598 -7.834 1.00 98.33 C
ATOM 20 C TYR A 3 1.232 2.440 -7.743 1.00 98.33 C
ATOM 21 CB TYR A 3 3.340 2.753 -6.431 1.00 98.33 C
ATOM 22 O TYR A 3 0.491 3.417 -7.877 1.00 98.33 O
ATOM 23 CG TYR A 3 4.849 2.744 -6.407 1.00 98.33 C
ATOM 24 CD1 TYR A 3 5.554 1.574 -6.131 1.00 98.33 C
ATOM 25 CD2 TYR A 3 5.573 3.904 -6.659 1.00 98.33 C
ATOM 26 CE1 TYR A 3 6.945 1.561 -6.106 1.00 98.33 C
ATOM 27 CE2 TYR A 3 6.964 3.903 -6.637 1.00 98.33 C
ATOM 28 OH TYR A 3 9.016 2.722 -6.336 1.00 98.33 O
ATOM 29 CZ TYR A 3 7.640 2.729 -6.360 1.00 98.33 C
ATOM 30 N LYS A 4 0.867 1.196 -7.558 1.00 98.45 N
ATOM 31 CA LYS A 4 -0.555 0.875 -7.488 1.00 98.45 C
ATOM 32 C LYS A 4 -0.916 0.269 -6.135 1.00 98.45 C
ATOM 33 CB LYS A 4 -0.946 -0.085 -8.613 1.00 98.45 C
ATOM 34 O LYS A 4 -0.130 -0.485 -5.557 1.00 98.45 O
ATOM 35 CG LYS A 4 -2.424 -0.444 -8.635 1.00 98.45 C
ATOM 36 CD LYS A 4 -2.728 -1.495 -9.695 1.00 98.45 C
ATOM 37 CE LYS A 4 -2.702 -0.901 -11.097 1.00 98.45 C
ATOM 38 NZ LYS A 4 -3.259 -1.846 -12.110 1.00 98.45 N
ATOM 39 N LEU A 5 -2.056 0.715 -5.667 1.00 98.19 N
ATOM 40 CA LEU A 5 -2.611 0.142 -4.446 1.00 98.19 C
ATOM 41 C LEU A 5 -3.944 -0.545 -4.724 1.00 98.19 C
ATOM 42 CB LEU A 5 -2.794 1.226 -3.380 1.00 98.19 C
ATOM 43 O LEU A 5 -4.835 0.046 -5.339 1.00 98.19 O
ATOM 44 CG LEU A 5 -3.556 0.815 -2.119 1.00 98.19 C
ATOM 45 CD1 LEU A 5 -2.758 -0.219 -1.331 1.00 98.19 C
ATOM 46 CD2 LEU A 5 -3.860 2.035 -1.256 1.00 98.19 C
ATOM 47 N ILE A 6 -4.005 -1.756 -4.246 1.00 97.96 N
ATOM 48 CA ILE A 6 -5.249 -2.516 -4.282 1.00 97.96 C
ATOM 49 C ILE A 6 -5.803 -2.666 -2.866 1.00 97.96 C
ATOM 50 CB ILE A 6 -5.045 -3.905 -4.928 1.00 97.96 C
ATOM 51 O ILE A 6 -5.137 -3.220 -1.989 1.00 97.96 O
ATOM 52 CG1 ILE A 6 -4.505 -3.756 -6.355 1.00 97.96 C
ATOM 53 CG2 ILE A 6 -6.352 -4.702 -4.917 1.00 97.96 C
ATOM 54 CD1 ILE A 6 -4.136 -5.076 -7.018 1.00 97.96 C
ATOM 55 N LEU A 7 -6.997 -2.109 -2.699 1.00 96.89 N
ATOM 56 CA LEU A 7 -7.681 -2.218 -1.415 1.00 96.89 C
ATOM 57 C LEU A 7 -8.701 -3.351 -1.437 1.00 96.89 C
ATOM 58 CB LEU A 7 -8.373 -0.899 -1.062 1.00 96.89 C
ATOM 59 O LEU A 7 -9.664 -3.310 -2.206 1.00 96.89 O
ATOM 60 CG LEU A 7 -7.460 0.309 -0.848 1.00 96.89 C
ATOM 61 CD1 LEU A 7 -8.177 1.593 -1.254 1.00 96.89 C
ATOM 62 CD2 LEU A 7 -7.003 0.383 0.605 1.00 96.89 C
ATOM 63 N ASN A 8 -8.362 -4.449 -0.693 1.00 93.97 N
ATOM 64 CA ASN A 8 -9.317 -5.537 -0.515 1.00 93.97 C
ATOM 65 C ASN A 8 -9.912 -5.537 0.891 1.00 93.97 C
ATOM 66 CB ASN A 8 -8.656 -6.885 -0.811 1.00 93.97 C
ATOM 67 O ASN A 8 -9.656 -6.450 1.678 1.00 93.97 O
ATOM 68 CG ASN A 8 -8.345 -7.072 -2.283 1.00 93.97 C
ATOM 69 ND2 ASN A 8 -7.265 -7.787 -2.575 1.00 93.97 N
ATOM 70 OD1 ASN A 8 -9.068 -6.576 -3.151 1.00 93.97 O
ATOM 71 N GLY A 9 -10.586 -4.520 1.150 1.00 87.95 N
ATOM 72 CA GLY A 9 -11.194 -4.470 2.470 1.00 87.95 C
ATOM 73 C GLY A 9 -12.660 -4.861 2.468 1.00 87.95 C
ATOM 74 O GLY A 9 -13.254 -5.055 1.405 1.00 87.95 O
ATOM 75 N LYS A 10 -13.204 -5.333 3.741 1.00 87.78 N
ATOM 76 CA LYS A 10 -14.627 -5.608 3.917 1.00 87.78 C
ATOM 77 C LYS A 10 -15.458 -4.341 3.738 1.00 87.78 C
ATOM 78 CB LYS A 10 -14.890 -6.217 5.295 1.00 87.78 C
ATOM 79 O LYS A 10 -16.536 -4.378 3.143 1.00 87.78 O
ATOM 80 CG LYS A 10 -14.226 -7.569 5.511 1.00 87.78 C
ATOM 81 CD LYS A 10 -14.567 -8.147 6.878 1.00 87.78 C
ATOM 82 CE LYS A 10 -13.892 -9.494 7.101 1.00 87.78 C
ATOM 83 NZ LYS A 10 -14.202 -10.052 8.451 1.00 87.78 N
ATOM 84 N THR A 11 -14.788 -3.122 4.178 1.00 85.58 N
ATOM 85 CA THR A 11 -15.487 -1.842 4.183 1.00 85.58 C
ATOM 86 C THR A 11 -15.236 -1.084 2.883 1.00 85.58 C
ATOM 87 CB THR A 11 -15.052 -0.973 5.378 1.00 85.58 C
ATOM 88 O THR A 11 -16.142 -0.440 2.349 1.00 85.58 O
ATOM 89 CG2 THR A 11 -15.583 -1.539 6.691 1.00 85.58 C
ATOM 90 OG1 THR A 11 -13.621 -0.928 5.434 1.00 85.58 O
ATOM 91 N LEU A 12 -14.051 -1.205 2.383 1.00 90.92 N
ATOM 92 CA LEU A 12 -13.651 -0.476 1.184 1.00 90.92 C
ATOM 93 C LEU A 12 -12.843 -1.369 0.249 1.00 90.92 C
ATOM 94 CB LEU A 12 -12.833 0.763 1.558 1.00 90.92 C
ATOM 95 O LEU A 12 -11.943 -2.086 0.692 1.00 90.92 O
ATOM 96 CG LEU A 12 -12.420 1.677 0.403 1.00 90.92 C
ATOM 97 CD1 LEU A 12 -13.644 2.360 -0.197 1.00 90.92 C
ATOM 98 CD2 LEU A 12 -11.403 2.710 0.876 1.00 90.92 C
ATOM 99 N LYS A 13 -13.275 -1.363 -1.001 1.00 94.92 N
ATOM 100 CA LYS A 13 -12.561 -2.059 -2.067 1.00 94.92 C
ATOM 101 C LYS A 13 -12.271 -1.124 -3.238 1.00 94.92 C
ATOM 102 CB LYS A 13 -13.363 -3.269 -2.549 1.00 94.92 C
ATOM 103 O LYS A 13 -13.066 -0.230 -3.536 1.00 94.92 O
ATOM 104 CG LYS A 13 -13.594 -4.325 -1.478 1.00 94.92 C
ATOM 105 CD LYS A 13 -14.368 -5.517 -2.027 1.00 94.92 C
ATOM 106 CE LYS A 13 -14.626 -6.561 -0.949 1.00 94.92 C
ATOM 107 NZ LYS A 13 -15.332 -7.759 -1.495 1.00 94.92 N
ATOM 108 N GLY A 14 -11.069 -1.258 -3.742 1.00 95.72 N
ATOM 109 CA GLY A 14 -10.753 -0.439 -4.902 1.00 95.72 C
ATOM 110 C GLY A 14 -9.267 -0.378 -5.202 1.00 95.72 C
ATOM 111 O GLY A 14 -8.481 -1.137 -4.630 1.00 95.72 O
ATOM 112 N GLU A 15 -9.041 0.411 -6.223 1.00 97.14 N
ATOM 113 CA GLU A 15 -7.661 0.598 -6.660 1.00 97.14 C
ATOM 114 C GLU A 15 -7.333 2.079 -6.833 1.00 97.14 C
ATOM 115 CB GLU A 15 -7.404 -0.154 -7.968 1.00 97.14 C
ATOM 116 O GLU A 15 -8.192 2.868 -7.231 1.00 97.14 O
ATOM 117 CG GLU A 15 -7.671 -1.650 -7.879 1.00 97.14 C
ATOM 118 CD GLU A 15 -7.530 -2.364 -9.214 1.00 97.14 C
ATOM 119 OE1 GLU A 15 -7.705 -3.603 -9.261 1.00 97.14 O
ATOM 120 OE2 GLU A 15 -7.240 -1.680 -10.221 1.00 97.14 O
ATOM 121 N THR A 16 -6.040 2.407 -6.429 1.00 96.97 N
ATOM 122 CA THR A 16 -5.551 3.759 -6.679 1.00 96.97 C
ATOM 123 C THR A 16 -4.062 3.743 -7.012 1.00 96.97 C
ATOM 124 CB THR A 16 -5.799 4.675 -5.466 1.00 96.97 C
ATOM 125 O THR A 16 -3.391 2.725 -6.831 1.00 96.97 O
ATOM 126 CG2 THR A 16 -4.916 4.278 -4.288 1.00 96.97 C
ATOM 127 OG1 THR A 16 -5.510 6.030 -5.831 1.00 96.97 O
ATOM 128 N THR A 17 -3.615 4.878 -7.608 1.00 98.25 N
ATOM 129 CA THR A 17 -2.208 4.959 -7.986 1.00 98.25 C
ATOM 130 C THR A 17 -1.563 6.213 -7.404 1.00 98.25 C
ATOM 131 CB THR A 17 -2.039 4.954 -9.516 1.00 98.25 C
ATOM 132 O THR A 17 -2.259 7.163 -7.039 1.00 98.25 O
ATOM 133 CG2 THR A 17 -2.603 3.675 -10.128 1.00 98.25 C
ATOM 134 OG1 THR A 17 -2.730 6.080 -10.071 1.00 98.25 O
ATOM 135 N THR A 18 -0.163 6.149 -7.219 1.00 98.09 N
ATOM 136 CA THR A 18 0.597 7.314 -6.780 1.00 98.09 C
ATOM 137 C THR A 18 1.978 7.335 -7.430 1.00 98.09 C
ATOM 138 CB THR A 18 0.748 7.338 -5.248 1.00 98.09 C
ATOM 139 O THR A 18 2.533 6.283 -7.755 1.00 98.09 O
ATOM 140 CG2 THR A 18 1.716 6.259 -4.774 1.00 98.09 C
ATOM 141 OG1 THR A 18 1.244 8.620 -4.842 1.00 98.09 O
ATOM 142 N GLU A 19 2.385 8.571 -7.640 1.00 97.84 N
ATOM 143 CA GLU A 19 3.765 8.730 -8.091 1.00 97.84 C
ATOM 144 C GLU A 19 4.714 8.915 -6.910 1.00 97.84 C
ATOM 145 CB GLU A 19 3.884 9.915 -9.053 1.00 97.84 C
ATOM 146 O GLU A 19 4.448 9.718 -6.014 1.00 97.84 O
ATOM 147 CG GLU A 19 3.177 9.701 -10.383 1.00 97.84 C
ATOM 148 CD GLU A 19 3.428 10.819 -11.382 1.00 97.84 C
ATOM 149 OE1 GLU A 19 3.374 10.562 -12.607 1.00 97.84 O
ATOM 150 OE2 GLU A 19 3.680 11.961 -10.937 1.00 97.84 O
ATOM 151 N ALA A 20 5.757 8.020 -6.889 1.00 97.35 N
ATOM 152 CA ALA A 20 6.745 8.109 -5.817 1.00 97.35 C
ATOM 153 C ALA A 20 8.140 7.758 -6.325 1.00 97.35 C
ATOM 154 CB ALA A 20 6.358 7.192 -4.659 1.00 97.35 C
ATOM 155 O ALA A 20 8.284 7.074 -7.341 1.00 97.35 O
ATOM 156 N VAL A 21 9.117 8.256 -5.592 1.00 96.93 N
ATOM 157 CA VAL A 21 10.505 8.060 -5.995 1.00 96.93 C
ATOM 158 C VAL A 21 10.934 6.626 -5.692 1.00 96.93 C
ATOM 159 CB VAL A 21 11.448 9.058 -5.287 1.00 96.93 C
ATOM 160 O VAL A 21 11.848 6.096 -6.329 1.00 96.93 O
ATOM 161 CG1 VAL A 21 11.160 10.488 -5.741 1.00 96.93 C
ATOM 162 CG2 VAL A 21 11.311 8.938 -3.771 1.00 96.93 C
ATOM 163 N ASP A 22 10.323 6.027 -4.673 1.00 97.54 N
ATOM 164 CA ASP A 22 10.642 4.654 -4.291 1.00 97.54 C
ATOM 165 C ASP A 22 9.463 3.995 -3.579 1.00 97.54 C
ATOM 166 CB ASP A 22 11.882 4.620 -3.396 1.00 97.54 C
ATOM 167 O ASP A 22 8.469 4.655 -3.271 1.00 97.54 O
ATOM 168 CG ASP A 22 11.745 5.488 -2.158 1.00 97.54 C
ATOM 169 OD1 ASP A 22 10.612 5.670 -1.663 1.00 97.54 O
ATOM 170 OD2 ASP A 22 12.779 5.998 -1.675 1.00 97.54 O
ATOM 171 N ALA A 23 9.542 2.683 -3.354 1.00 96.84 N
ATOM 172 CA ALA A 23 8.463 1.899 -2.759 1.00 96.84 C
ATOM 173 C ALA A 23 8.178 2.354 -1.331 1.00 96.84 C
ATOM 174 CB ALA A 23 8.810 0.412 -2.782 1.00 96.84 C
ATOM 175 O ALA A 23 7.021 2.396 -0.905 1.00 96.84 O
ATOM 176 N ALA A 24 9.209 2.742 -0.646 1.00 96.88 N
ATOM 177 CA ALA A 24 9.044 3.171 0.740 1.00 96.88 C
ATOM 178 C ALA A 24 8.183 4.429 0.826 1.00 96.88 C
ATOM 179 CB ALA A 24 10.404 3.417 1.388 1.00 96.88 C
ATOM 180 O ALA A 24 7.310 4.533 1.692 1.00 96.88 O
ATOM 181 N THR A 25 8.439 5.396 -0.120 1.00 97.59 N
ATOM 182 CA THR A 25 7.652 6.623 -0.168 1.00 97.59 C
ATOM 183 C THR A 25 6.212 6.327 -0.578 1.00 97.59 C
ATOM 184 CB THR A 25 8.266 7.644 -1.143 1.00 97.59 C
ATOM 185 O THR A 25 5.271 6.884 -0.007 1.00 97.59 O
ATOM 186 CG2 THR A 25 7.486 8.954 -1.133 1.00 97.59 C
ATOM 187 OG1 THR A 25 9.622 7.905 -0.761 1.00 97.59 O
ATOM 188 N ALA A 26 6.033 5.422 -1.520 1.00 97.91 N
ATOM 189 CA ALA A 26 4.704 5.005 -1.961 1.00 97.91 C
ATOM 190 C ALA A 26 3.945 4.311 -0.834 1.00 97.91 C
ATOM 191 CB ALA A 26 4.811 4.083 -3.173 1.00 97.91 C
ATOM 192 O ALA A 26 2.751 4.552 -0.641 1.00 97.91 O
ATOM 193 N GLU A 27 4.605 3.513 -0.093 1.00 97.35 N
ATOM 194 CA GLU A 27 3.990 2.783 1.011 1.00 97.35 C
ATOM 195 C GLU A 27 3.416 3.739 2.053 1.00 97.35 C
ATOM 196 CB GLU A 27 5.003 1.840 1.664 1.00 97.35 C
ATOM 197 O GLU A 27 2.303 3.537 2.542 1.00 97.35 O
ATOM 198 CG GLU A 27 4.423 0.998 2.791 1.00 97.35 C
ATOM 199 CD GLU A 27 5.423 0.019 3.385 1.00 97.35 C
ATOM 200 OE1 GLU A 27 5.038 -0.781 4.268 1.00 97.35 O
ATOM 201 OE2 GLU A 27 6.601 0.052 2.964 1.00 97.35 O
ATOM 202 N LYS A 28 4.154 4.789 2.410 1.00 97.42 N
ATOM 203 CA LYS A 28 3.692 5.772 3.386 1.00 97.42 C
ATOM 204 C LYS A 28 2.416 6.461 2.909 1.00 97.42 C
ATOM 205 CB LYS A 28 4.779 6.812 3.657 1.00 97.42 C
ATOM 206 O LYS A 28 1.484 6.663 3.690 1.00 97.42 O
ATOM 207 CG LYS A 28 5.958 6.281 4.460 1.00 97.42 C
ATOM 208 CD LYS A 28 6.975 7.377 4.752 1.00 97.42 C
ATOM 209 CE LYS A 28 8.177 6.836 5.513 1.00 97.42 C
ATOM 210 NZ LYS A 28 9.230 7.879 5.699 1.00 97.42 N
ATOM 211 N VAL A 29 2.381 6.745 1.649 1.00 97.05 N
ATOM 212 CA VAL A 29 1.226 7.414 1.059 1.00 97.05 C
ATOM 213 C VAL A 29 0.017 6.483 1.093 1.00 97.05 C
ATOM 214 CB VAL A 29 1.511 7.866 -0.391 1.00 97.05 C
ATOM 215 O VAL A 29 -1.072 6.885 1.511 1.00 97.05 O
ATOM 216 CG1 VAL A 29 0.240 8.402 -1.048 1.00 97.05 C
ATOM 217 CG2 VAL A 29 2.614 8.922 -0.414 1.00 97.05 C
ATOM 218 N PHE A 30 0.261 5.260 0.714 1.00 97.61 N
ATOM 219 CA PHE A 30 -0.831 4.296 0.647 1.00 97.61 C
ATOM 220 C PHE A 30 -1.343 3.960 2.042 1.00 97.61 C
ATOM 221 CB PHE A 30 -0.379 3.019 -0.069 1.00 97.61 C
ATOM 222 O PHE A 30 -2.552 3.846 2.255 1.00 97.61 O
ATOM 223 CG PHE A 30 -0.171 3.194 -1.549 1.00 97.61 C
ATOM 224 CD1 PHE A 30 -0.974 4.060 -2.280 1.00 97.61 C
ATOM 225 CD2 PHE A 30 0.829 2.491 -2.209 1.00 97.61 C
ATOM 226 CE1 PHE A 30 -0.784 4.223 -3.651 1.00 97.61 C
ATOM 227 CE2 PHE A 30 1.025 2.649 -3.578 1.00 97.61 C
ATOM 228 CZ PHE A 30 0.217 3.515 -4.297 1.00 97.61 C
ATOM 229 N LYS A 31 -0.477 3.819 3.011 1.00 96.04 N
ATOM 230 CA LYS A 31 -0.896 3.521 4.377 1.00 96.04 C
ATOM 231 C LYS A 31 -1.736 4.656 4.955 1.00 96.04 C
ATOM 232 CB LYS A 31 0.321 3.264 5.268 1.00 96.04 C
ATOM 233 O LYS A 31 -2.725 4.413 5.649 1.00 96.04 O
ATOM 234 CG LYS A 31 0.950 1.892 5.075 1.00 96.04 C
ATOM 235 CD LYS A 31 2.047 1.631 6.099 1.00 96.04 C
ATOM 236 CE LYS A 31 2.551 0.196 6.027 1.00 96.04 C
ATOM 237 NZ LYS A 31 3.617 -0.069 7.040 1.00 96.04 N
ATOM 238 N GLN A 32 -1.333 5.832 4.614 1.00 96.45 N
ATOM 239 CA GLN A 32 -2.122 6.977 5.057 1.00 96.45 C
ATOM 240 C GLN A 32 -3.501 6.980 4.404 1.00 96.45 C
ATOM 241 CB GLN A 32 -1.391 8.284 4.747 1.00 96.45 C
ATOM 242 O GLN A 32 -4.508 7.230 5.070 1.00 96.45 O
ATOM 243 CG GLN A 32 -2.088 9.522 5.296 1.00 96.45 C
ATOM 244 CD GLN A 32 -2.164 9.530 6.812 1.00 96.45 C
ATOM 245 NE2 GLN A 32 -3.376 9.640 7.344 1.00 96.45 N
ATOM 246 OE1 GLN A 32 -1.141 9.436 7.498 1.00 96.45 O
ATOM 247 N TYR A 33 -3.483 6.750 3.178 1.00 95.33 N
ATOM 248 CA TYR A 33 -4.737 6.673 2.435 1.00 95.33 C
ATOM 249 C TYR A 33 -5.650 5.600 3.017 1.00 95.33 C
ATOM 250 CB TYR A 33 -4.468 6.384 0.955 1.00 95.33 C
ATOM 251 O TYR A 33 -6.839 5.841 3.238 1.00 95.33 O
ATOM 252 CG TYR A 33 -5.722 6.263 0.124 1.00 95.33 C
ATOM 253 CD1 TYR A 33 -6.260 5.016 -0.183 1.00 95.33 C
ATOM 254 CD2 TYR A 33 -6.370 7.396 -0.358 1.00 95.33 C
ATOM 255 CE1 TYR A 33 -7.414 4.899 -0.951 1.00 95.33 C
ATOM 256 CE2 TYR A 33 -7.525 7.291 -1.126 1.00 95.33 C
ATOM 257 OH TYR A 33 -9.181 5.931 -2.178 1.00 95.33 O
ATOM 258 CZ TYR A 33 -8.039 6.040 -1.417 1.00 95.33 C
ATOM 259 N ALA A 34 -5.152 4.410 3.257 1.00 94.87 N
ATOM 260 CA ALA A 34 -5.929 3.318 3.838 1.00 94.87 C
ATOM 261 C ALA A 34 -6.432 3.684 5.232 1.00 94.87 C
ATOM 262 CB ALA A 34 -5.094 2.042 3.895 1.00 94.87 C
ATOM 263 O ALA A 34 -7.601 3.460 5.556 1.00 94.87 O
ATOM 264 N ASN A 35 -5.598 4.305 6.007 1.00 93.98 N
ATOM 265 CA ASN A 35 -5.967 4.719 7.356 1.00 93.98 C
ATOM 266 C ASN A 35 -7.077 5.766 7.338 1.00 93.98 C
ATOM 267 CB ASN A 35 -4.745 5.254 8.107 1.00 93.98 C
ATOM 268 O ASN A 35 -8.024 5.687 8.122 1.00 93.98 O
ATOM 269 CG ASN A 35 -5.028 5.513 9.574 1.00 93.98 C
ATOM 270 ND2 ASN A 35 -4.836 4.495 10.403 1.00 93.98 N
ATOM 271 OD1 ASN A 35 -5.417 6.619 9.957 1.00 93.98 O
ATOM 272 N ASP A 36 -6.987 6.701 6.417 1.00 95.00 N
ATOM 273 CA ASP A 36 -7.985 7.756 6.277 1.00 95.00 C
ATOM 274 C ASP A 36 -9.344 7.179 5.888 1.00 95.00 C
ATOM 275 CB ASP A 36 -7.535 8.786 5.239 1.00 95.00 C
ATOM 276 O ASP A 36 -10.386 7.737 6.240 1.00 95.00 O
ATOM 277 CG ASP A 36 -6.422 9.689 5.743 1.00 95.00 C
ATOM 278 OD1 ASP A 36 -6.138 9.685 6.960 1.00 95.00 O
ATOM 279 OD2 ASP A 36 -5.827 10.413 4.916 1.00 95.00 O
ATOM 280 N ASN A 37 -9.207 5.989 5.319 1.00 92.77 N
ATOM 281 CA ASN A 37 -10.449 5.372 4.866 1.00 92.77 C
ATOM 282 C ASN A 37 -10.870 4.223 5.778 1.00 92.77 C
ATOM 283 CB ASN A 37 -10.310 4.882 3.423 1.00 92.77 C
ATOM 284 O ASN A 37 -11.794 3.475 5.454 1.00 92.77 O
ATOM 285 CG ASN A 37 -10.384 6.010 2.413 1.00 92.77 C
ATOM 286 ND2 ASN A 37 -9.228 6.452 1.932 1.00 92.77 N
ATOM 287 OD1 ASN A 37 -11.471 6.481 2.069 1.00 92.77 O
ATOM 288 N GLY A 38 -10.171 4.142 6.859 1.00 92.50 N
ATOM 289 CA GLY A 38 -10.543 3.168 7.873 1.00 92.50 C
ATOM 290 C GLY A 38 -10.153 1.748 7.508 1.00 92.50 C
ATOM 291 O GLY A 38 -10.788 0.791 7.955 1.00 92.50 O
ATOM 292 N VAL A 39 -9.272 1.633 6.574 1.00 93.16 N
ATOM 293 CA VAL A 39 -8.765 0.326 6.168 1.00 93.16 C
ATOM 294 C VAL A 39 -7.482 0.009 6.934 1.00 93.16 C
ATOM 295 CB VAL A 39 -8.506 0.266 4.646 1.00 93.16 C
ATOM 296 O VAL A 39 -6.502 0.753 6.851 1.00 93.16 O
ATOM 297 CG1 VAL A 39 -7.921 -1.089 4.251 1.00 93.16 C
ATOM 298 CG2 VAL A 39 -9.795 0.541 3.875 1.00 93.16 C
ATOM 299 N ASP A 40 -7.638 -1.000 7.794 1.00 91.01 N
ATOM 300 CA ASP A 40 -6.484 -1.557 8.493 1.00 91.01 C
ATOM 301 C ASP A 40 -6.344 -3.052 8.217 1.00 91.01 C
ATOM 302 CB ASP A 40 -6.596 -1.307 9.998 1.00 91.01 C
ATOM 303 O ASP A 40 -7.224 -3.839 8.573 1.00 91.01 O
ATOM 304 CG ASP A 40 -5.345 -1.710 10.760 1.00 91.01 C
ATOM 305 OD1 ASP A 40 -4.285 -1.911 10.130 1.00 91.01 O
ATOM 306 OD2 ASP A 40 -5.422 -1.830 12.002 1.00 91.01 O
ATOM 307 N GLY A 41 -5.303 -3.437 7.548 1.00 92.21 N
ATOM 308 CA GLY A 41 -5.149 -4.844 7.216 1.00 92.21 C
ATOM 309 C GLY A 41 -3.713 -5.234 6.917 1.00 92.21 C
ATOM 310 O GLY A 41 -2.779 -4.550 7.340 1.00 92.21 O
ATOM 311 N GLU A 42 -3.628 -6.454 6.389 1.00 95.96 N
ATOM 312 CA GLU A 42 -2.327 -6.995 6.004 1.00 95.96 C
ATOM 313 C GLU A 42 -1.863 -6.418 4.670 1.00 95.96 C
ATOM 314 CB GLU A 42 -2.381 -8.523 5.927 1.00 95.96 C
ATOM 315 O GLU A 42 -2.633 -6.370 3.708 1.00 95.96 O
ATOM 316 CG GLU A 42 -2.443 -9.205 7.286 1.00 95.96 C
ATOM 317 CD GLU A 42 -2.186 -10.702 7.218 1.00 95.96 C
ATOM 318 OE1 GLU A 42 -1.205 -11.178 7.834 1.00 95.96 O
ATOM 319 OE2 GLU A 42 -2.972 -11.404 6.543 1.00 95.96 O
ATOM 320 N TRP A 43 -0.525 -6.041 4.695 1.00 96.56 N
ATOM 321 CA TRP A 43 0.030 -5.392 3.512 1.00 96.56 C
ATOM 322 C TRP A 43 1.004 -6.317 2.791 1.00 96.56 C
ATOM 323 CB TRP A 43 0.733 -4.086 3.894 1.00 96.56 C
ATOM 324 O TRP A 43 1.795 -7.015 3.430 1.00 96.56 O
ATOM 325 CG TRP A 43 -0.204 -2.978 4.269 1.00 96.56 C
ATOM 326 CD1 TRP A 43 -0.857 -2.818 5.459 1.00 96.56 C
ATOM 327 CD2 TRP A 43 -0.596 -1.876 3.445 1.00 96.56 C
ATOM 328 CE2 TRP A 43 -1.490 -1.085 4.201 1.00 96.56 C
ATOM 329 CE3 TRP A 43 -0.278 -1.479 2.139 1.00 96.56 C
ATOM 330 NE1 TRP A 43 -1.631 -1.681 5.425 1.00 96.56 N
ATOM 331 CH2 TRP A 43 -1.741 0.446 2.414 1.00 96.56 C
ATOM 332 CZ2 TRP A 43 -2.069 0.081 3.694 1.00 96.56 C
ATOM 333 CZ3 TRP A 43 -0.856 -0.319 1.636 1.00 96.56 C
ATOM 334 N THR A 44 0.864 -6.326 1.450 1.00 97.82 N
ATOM 335 CA THR A 44 1.816 -7.030 0.598 1.00 97.82 C
ATOM 336 C THR A 44 2.276 -6.137 -0.551 1.00 97.82 C
ATOM 337 CB THR A 44 1.208 -8.326 0.032 1.00 97.82 C
ATOM 338 O THR A 44 1.544 -5.242 -0.979 1.00 97.82 O
ATOM 339 CG2 THR A 44 0.804 -9.279 1.153 1.00 97.82 C
ATOM 340 OG1 THR A 44 0.049 -8.004 -0.746 1.00 97.82 O
ATOM 341 N TYR A 45 3.536 -6.362 -0.850 1.00 97.93 N
ATOM 342 CA TYR A 45 4.082 -5.581 -1.955 1.00 97.93 C
ATOM 343 C TYR A 45 4.717 -6.487 -3.002 1.00 97.93 C
ATOM 344 CB TYR A 45 5.114 -4.572 -1.442 1.00 97.93 C
ATOM 345 O TYR A 45 5.526 -7.357 -2.671 1.00 97.93 O
ATOM 346 CG TYR A 45 5.833 -3.828 -2.541 1.00 97.93 C
ATOM 347 CD1 TYR A 45 7.215 -3.925 -2.686 1.00 97.93 C
ATOM 348 CD2 TYR A 45 5.133 -3.026 -3.436 1.00 97.93 C
ATOM 349 CE1 TYR A 45 7.882 -3.239 -3.696 1.00 97.93 C
ATOM 350 CE2 TYR A 45 5.790 -2.336 -4.449 1.00 97.93 C
ATOM 351 OH TYR A 45 7.818 -1.768 -5.573 1.00 97.93 O
ATOM 352 CZ TYR A 45 7.162 -2.449 -4.571 1.00 97.93 C
ATOM 353 N ASP A 46 4.233 -6.194 -4.296 1.00 97.90 N
ATOM 354 CA ASP A 46 4.795 -6.866 -5.464 1.00 97.90 C
ATOM 355 C ASP A 46 5.693 -5.922 -6.260 1.00 97.90 C
ATOM 356 CB ASP A 46 3.680 -7.411 -6.358 1.00 97.90 C
ATOM 357 O ASP A 46 5.204 -5.027 -6.953 1.00 97.90 O
ATOM 358 CG ASP A 46 4.200 -8.272 -7.496 1.00 97.90 C
ATOM 359 OD1 ASP A 46 5.374 -8.113 -7.894 1.00 97.90 O
ATOM 360 OD2 ASP A 46 3.427 -9.115 -8.001 1.00 97.90 O
ATOM 361 N ASP A 47 7.035 -6.192 -6.200 1.00 96.38 N
ATOM 362 CA ASP A 47 8.016 -5.286 -6.790 1.00 96.38 C
ATOM 363 C ASP A 47 7.993 -5.370 -8.315 1.00 96.38 C
ATOM 364 CB ASP A 47 9.419 -5.599 -6.267 1.00 96.38 C
ATOM 365 O ASP A 47 8.258 -4.381 -9.000 1.00 96.38 O
ATOM 366 CG ASP A 47 10.456 -4.581 -6.706 1.00 96.38 C
ATOM 367 OD1 ASP A 47 10.283 -3.374 -6.428 1.00 96.38 O
ATOM 368 OD2 ASP A 47 11.455 -4.988 -7.337 1.00 96.38 O
ATOM 369 N ALA A 48 7.591 -6.528 -8.836 1.00 96.52 N
ATOM 370 CA ALA A 48 7.536 -6.733 -10.281 1.00 96.52 C
ATOM 371 C ALA A 48 6.471 -5.848 -10.921 1.00 96.52 C
ATOM 372 CB ALA A 48 7.266 -8.202 -10.601 1.00 96.52 C
ATOM 373 O ALA A 48 6.701 -5.253 -11.976 1.00 96.52 O
ATOM 374 N THR A 49 5.425 -5.657 -10.194 1.00 97.33 N
ATOM 375 CA THR A 49 4.307 -4.909 -10.757 1.00 97.33 C
ATOM 376 C THR A 49 4.128 -3.577 -10.032 1.00 97.33 C
ATOM 377 CB THR A 49 2.999 -5.717 -10.680 1.00 97.33 C
ATOM 378 O THR A 49 3.210 -2.816 -10.341 1.00 97.33 O
ATOM 379 CG2 THR A 49 3.124 -7.038 -11.432 1.00 97.33 C
ATOM 380 OG1 THR A 49 2.691 -5.989 -9.307 1.00 97.33 O
ATOM 381 N LYS A 50 4.988 -3.309 -9.054 1.00 97.22 N
ATOM 382 CA LYS A 50 4.913 -2.103 -8.233 1.00 97.22 C
ATOM 383 C LYS A 50 3.519 -1.928 -7.639 1.00 97.22 C
ATOM 384 CB LYS A 50 5.292 -0.870 -9.055 1.00 97.22 C
ATOM 385 O LYS A 50 2.962 -0.828 -7.663 1.00 97.22 O
ATOM 386 CG LYS A 50 6.668 -0.954 -9.700 1.00 97.22 C
ATOM 387 CD LYS A 50 7.774 -0.986 -8.654 1.00 97.22 C
ATOM 388 CE LYS A 50 9.154 -0.933 -9.296 1.00 97.22 C
ATOM 389 NZ LYS A 50 10.016 -2.069 -8.853 1.00 97.22 N
ATOM 390 N THR A 51 2.970 -3.067 -7.091 1.00 98.49 N
ATOM 391 CA THR A 51 1.600 -3.090 -6.593 1.00 98.49 C
ATOM 392 C THR A 51 1.569 -3.442 -5.109 1.00 98.49 C
ATOM 393 CB THR A 51 0.737 -4.095 -7.380 1.00 98.49 C
ATOM 394 O THR A 51 2.172 -4.432 -4.688 1.00 98.49 O
ATOM 395 CG2 THR A 51 -0.695 -4.120 -6.855 1.00 98.49 C
ATOM 396 OG1 THR A 51 0.721 -3.721 -8.763 1.00 98.49 O
ATOM 397 N PHE A 52 0.876 -2.509 -4.382 1.00 98.30 N
ATOM 398 CA PHE A 52 0.561 -2.767 -2.982 1.00 98.30 C
ATOM 399 C PHE A 52 -0.853 -3.319 -2.839 1.00 98.30 C
ATOM 400 CB PHE A 52 0.713 -1.490 -2.151 1.00 98.30 C
ATOM 401 O PHE A 52 -1.770 -2.879 -3.535 1.00 98.30 O
ATOM 402 CG PHE A 52 2.138 -1.033 -1.995 1.00 98.30 C
ATOM 403 CD1 PHE A 52 2.876 -1.383 -0.871 1.00 98.30 C
ATOM 404 CD2 PHE A 52 2.741 -0.253 -2.974 1.00 98.30 C
ATOM 405 CE1 PHE A 52 4.196 -0.962 -0.725 1.00 98.30 C
ATOM 406 CE2 PHE A 52 4.059 0.172 -2.834 1.00 98.30 C
ATOM 407 CZ PHE A 52 4.784 -0.183 -1.709 1.00 98.30 C
ATOM 408 N THR A 53 -0.919 -4.370 -1.938 1.00 98.18 N
ATOM 409 CA THR A 53 -2.233 -4.934 -1.648 1.00 98.18 C
ATOM 410 C THR A 53 -2.497 -4.946 -0.145 1.00 98.18 C
ATOM 411 CB THR A 53 -2.362 -6.364 -2.206 1.00 98.18 C
ATOM 412 O THR A 53 -1.641 -5.364 0.637 1.00 98.18 O
ATOM 413 CG2 THR A 53 -3.762 -6.921 -1.973 1.00 98.18 C
ATOM 414 OG1 THR A 53 -2.095 -6.347 -3.614 1.00 98.18 O
ATOM 415 N VAL A 54 -3.608 -4.311 0.169 1.00 97.01 N
ATOM 416 CA VAL A 54 -4.037 -4.394 1.561 1.00 97.01 C
ATOM 417 C VAL A 54 -5.355 -5.160 1.651 1.00 97.01 C
ATOM 418 CB VAL A 54 -4.192 -2.991 2.191 1.00 97.01 C
ATOM 419 O VAL A 54 -6.284 -4.903 0.881 1.00 97.01 O
ATOM 420 CG1 VAL A 54 -5.209 -2.159 1.413 1.00 97.01 C
ATOM 421 CG2 VAL A 54 -4.603 -3.108 3.658 1.00 97.01 C
ATOM 422 N THR A 55 -5.311 -6.195 2.592 1.00 95.82 N
ATOM 423 CA THR A 55 -6.495 -7.009 2.844 1.00 95.82 C
ATOM 424 C THR A 55 -6.937 -6.885 4.299 1.00 95.82 C
ATOM 425 CB THR A 55 -6.238 -8.489 2.507 1.00 95.82 C
ATOM 426 O THR A 55 -6.134 -7.074 5.216 1.00 95.82 O
ATOM 427 CG2 THR A 55 -7.514 -9.314 2.642 1.00 95.82 C
ATOM 428 OG1 THR A 55 -5.757 -8.590 1.161 1.00 95.82 O
ATOM 429 N GLU A 56 -8.163 -6.316 4.445 1.00 92.02 N
ATOM 430 CA GLU A 56 -8.777 -6.262 5.768 1.00 92.02 C
ATOM 431 C GLU A 56 -9.401 -7.603 6.143 1.00 92.02 C
ATOM 432 CB GLU A 56 -9.834 -5.156 5.825 1.00 92.02 C
ATOM 433 O GLU A 56 -9.881 -8.334 5.274 1.00 92.02 O
ATOM 434 CG GLU A 56 -10.379 -4.898 7.223 1.00 92.02 C
ATOM 435 CD GLU A 56 -11.456 -3.826 7.258 1.00 92.02 C
ATOM 436 OE1 GLU A 56 -12.006 -3.553 8.349 1.00 92.02 O
ATOM 437 OE2 GLU A 56 -11.754 -3.255 6.185 1.00 92.02 O
TER 438 GLU A 56
ENDMDL
END

Просмотреть файл

@ -3,30 +3,63 @@ import torch.nn as nn
from sequence_models.constants import PROTEIN_ALPHABET, PAD
from sequence_models.convolutional import ByteNetLM
from sequence_models.collaters import SimpleCollater
from sequence_models.gnn import BidirectionalStruct2SeqDecoder
from sequence_models.collaters import SimpleCollater, StructureCollater
CARP_URL = 'https://zenodo.org/record/6564798/files/'
MIF_URL = 'https://zenodo.org/record/6573779/files/'
n_tokens = len(PROTEIN_ALPHABET)
def load_carp(model_data):
d_embedding = model_data['d_embed']
d_model = model_data['d_model']
n_layers = model_data['n_layers']
kernel_size = model_data['kernel_size']
activation = model_data['activation']
slim = model_data['slim']
r = model_data['r']
model = ByteNetLM(n_tokens, d_embedding, d_model, n_layers, kernel_size, r, dropout=0.0,
activation=activation, causal=False, padding_idx=PROTEIN_ALPHABET.index(PAD),
final_ln=True, slim=slim)
sd = model_data['model_state_dict']
model.load_state_dict(sd)
model = CARP(model)
return model
def load_gnn(model_data):
one_hot_src = model_data['model'] == 'mif'
gnn = BidirectionalStruct2SeqDecoder(n_tokens, 10, 11,
256, num_decoder_layers=4,
dropout=0.05, use_mpnn=True,
pe=False, one_hot_src=one_hot_src)
sd = model_data['model_state_dict']
gnn.load_state_dict(sd)
return gnn
def load_model_and_alphabet(model_name):
if not model_name.endswith(".pt"): # treat as filepath
url = 'https://zenodo.org/record/6368484/files/%s.pt?download=1' %model_name
model_data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
if 'carp' in model_name:
url = CARP_URL + '%s.pt?download=1' %model_name
model_data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
elif 'mif' in model_name:
url = MIF_URL + '%s.pt?download=1' %model_name
model_data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
else:
model_data = torch.load(model_name, map_location="cpu")
sd = model_data['model_state_dict']
n_tokens = len(PROTEIN_ALPHABET)
collater = SimpleCollater(PROTEIN_ALPHABET, pad=True)
if model_data['model'] == 'carp':
d_embedding = model_data['d_embed']
d_model = model_data['d_model']
n_layers = model_data['n_layers']
kernel_size = model_data['kernel_size']
activation = model_data['activation']
slim = model_data['slim']
r = model_data['r']
model = ByteNetLM(n_tokens, d_embedding, d_model, n_layers, kernel_size, r, dropout=0.0,
activation=activation, causal=False, padding_idx=PROTEIN_ALPHABET.index(PAD),
final_ln=True, slim=slim)
collater = SimpleCollater(PROTEIN_ALPHABET, pad=True)
model.load_state_dict(sd)
model = load_carp(model_data)
elif model_data['model'] in ['mif', 'mif-st']:
gnn = load_gnn(model_data)
cnn = None
if model_data['model'] == 'mif-st':
url = CARP_URL + '%s.pt?download=1' % 'carp_640M'
cnn_data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
cnn = load_carp(cnn_data)
collater = StructureCollater(collater, n_connections=30)
model = MIF(gnn, cnn=cnn)
return model, collater
@ -38,7 +71,8 @@ class CARP(nn.Module):
self.model = model
def forward(self, x, result='repr'):
padding_mask = x == PROTEIN_ALPHABET.index(PAD)
padding_mask = (x == PROTEIN_ALPHABET.index(PAD))
padding_mask = padding_mask.unsqueeze(-1)
if result == 'repr':
return self.model.embedder(x, input_mask=padding_mask)
elif result == 'logits':
@ -46,3 +80,23 @@ class CARP(nn.Module):
else:
raise ValueError("Result must be either 'repr' or 'logits'")
class MIF(nn.Module):
"""Wrapper that takes care of input masking."""
def __init__(self, gnn: BidirectionalStruct2SeqDecoder, cnn=None):
super().__init__()
self.gnn = gnn
self.cnn = cnn
def forward(self, src, nodes, edges, connections, edge_mask, result='repr'):
if result == 'logits':
decoder = True
elif result == 'repr':
decoder = False
else:
raise ValueError("Result must be either 'repr' or 'logits'")
if self.cnn is not None:
src = self.cnn(src, result='logits')
return self.gnn(nodes, edges, connections, src, edge_mask, decoder=decoder)

Просмотреть файл

@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
setuptools.setup(
name="sequence-models",
version="1.0.0",
version="1.1.0",
author="Kevin Yang",
author_email="yang.kevin@microsoft.com",
description="Machine learning for sequences.",