habdine commited on
Commit
acb7e86
1 Parent(s): 0d91f43

Upload code

Browse files
Files changed (8) hide show
  1. configuration_prot2text.py +74 -0
  2. conversion.py +470 -0
  3. graphs.py +1137 -0
  4. modeling_prot2text.py +392 -0
  5. pdb2graph.py +171 -0
  6. utils.py +742 -0
  7. utils_convert.py +82 -0
  8. utils_dataset.py +60 -0
configuration_prot2text.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ Prot2Text configuration"""
2
+
3
+ from transformers.configuration_utils import PretrainedConfig
4
+ from transformers import AutoConfig
5
+ from transformers.utils import logging
6
+
7
+
8
+ logger = logging.get_logger(__name__)
9
+
10
+
11
+ class Prot2TextConfig(PretrainedConfig):
12
+ model_type = "prot2text"
13
+ keys_to_ignore_at_inference = ["past_key_values"]
14
+ _keys_to_ignore_on_load_missing = [r"transformer"]
15
+
16
+ def __init__(
17
+ self,
18
+ cross_esm_graph=True,
19
+ decoder_start_token_id=50257,
20
+ early_stopping=True,
21
+ eos_token_id=50258,
22
+ bos_token_id=50257,
23
+ esm=True,
24
+ esm_model_name="facebook/esm2_t6_8M_UR50D",
25
+ gpt_model_name="gpt2",
26
+ length_penalty=2.0,
27
+ max_new_tokens=256,
28
+ no_repeat_ngram_size=3,
29
+ pad_token_id=50256,
30
+ prot2text_version="1.1",
31
+ rgcn=True,
32
+ rgc_input_dim=67,
33
+ rgcn_n_layers=6,
34
+ gpt_config=None,
35
+ esm_config=None,
36
+ **kwargs,
37
+ ):
38
+ self.cross_esm_graph = cross_esm_graph
39
+ self.decoder_start_token_id = decoder_start_token_id
40
+ self.early_stopping = early_stopping
41
+ self.eos_token_id = eos_token_id
42
+ self.esm = esm
43
+ self.esm_model_name = esm_model_name
44
+ self.gpt_model_name = gpt_model_name
45
+ self.length_penalty = length_penalty
46
+ self.max_new_tokens = max_new_tokens
47
+ self.no_repeat_ngram_size = no_repeat_ngram_size
48
+ self.pad_token_id = pad_token_id
49
+ self.prot2text_version = prot2text_version
50
+ self.rgcn = rgcn
51
+ self.rgc_input_dim = rgc_input_dim
52
+ self.rgcn_n_layers = rgcn_n_layers
53
+ if gpt_config is None:
54
+ self.gpt_config = AutoConfig.from_pretrained(gpt_model_name,
55
+ _name_or_path= gpt_model_name,
56
+ is_encoder_decoder=True,
57
+ use_cache=False,
58
+ add_cross_attention=True,
59
+ bos_token_id=bos_token_id,
60
+ decoder_start_token_id=decoder_start_token_id,
61
+ eos_token_id=eos_token_id,
62
+ max_new_tokens=max_new_tokens,
63
+ pad_token_id=50256,
64
+ vocab_size=50259,
65
+ num_beams=1,
66
+ max_length=256,
67
+ min_length=1).to_dict()
68
+ else:
69
+ self.gpt_config = gpt_config
70
+ if esm_config is None:
71
+ self.esm_config = AutoConfig.from_pretrained(esm_model_name).to_dict()
72
+ self.esm_config = esm_config
73
+
74
+ super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
conversion.py ADDED
@@ -0,0 +1,470 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Utilities for converting Graphein Networks to Geometric Deep Learning formats.
2
+ """
3
+ # %%
4
+ # Graphein
5
+ # Author: Kexin Huang, Arian Jamasb <[email protected]>
6
+ # License: MIT
7
+ # Project Website: https://github.com/a-r-j/graphein
8
+ # Code Repository: https://github.com/a-r-j/graphein
9
+ from __future__ import annotations
10
+
11
+ from typing import List, Optional
12
+
13
+ import networkx as nx
14
+ import numpy as np
15
+ import torch
16
+
17
+ from graphein.utils.dependencies import import_message
18
+
19
+ try:
20
+ import torch_geometric
21
+ from torch_geometric.data import Data
22
+ except ImportError:
23
+ import_message(
24
+ submodule="graphein.ml.conversion",
25
+ package="torch_geometric",
26
+ pip_install=True,
27
+ conda_channel="rusty1s",
28
+ )
29
+
30
+ try:
31
+ import dgl
32
+ except ImportError:
33
+ import_message(
34
+ submodule="graphein.ml.conversion",
35
+ package="dgl",
36
+ pip_install=True,
37
+ conda_channel="dglteam",
38
+ )
39
+
40
+ try:
41
+ import jax.numpy as jnp
42
+ except ImportError:
43
+ import_message(
44
+ submodule="graphein.ml.conversion",
45
+ package="jax",
46
+ pip_install=True,
47
+ conda_channel="conda-forge",
48
+ )
49
+ try:
50
+ import jraph
51
+ except ImportError:
52
+ import_message(
53
+ submodule="graphein.ml.conversion",
54
+ package="jraph",
55
+ pip_install=True,
56
+ conda_channel="conda-forge",
57
+ )
58
+
59
+
60
+ SUPPORTED_FORMATS = ["nx", "pyg", "dgl", "jraph"]
61
+ """Supported conversion formats.
62
+
63
+ ``"nx"``: NetworkX graph
64
+
65
+ ``"pyg"``: PyTorch Geometric Data object
66
+
67
+ ``"dgl"``: DGL graph
68
+
69
+ ``"Jraph"``: Jraph GraphsTuple
70
+ """
71
+
72
+ SUPPORTED_VERBOSITY = ["gnn", "default", "all_info"]
73
+ """Supported verbosity levels for preserving graph features in conversion."""
74
+
75
+
76
+ class GraphFormatConvertor:
77
+ """
78
+ Provides conversion utilities between NetworkX Graphs and geometric deep learning library destination formats.
79
+ Currently, we provide support for converstion from ``nx.Graph`` to ``dgl.DGLGraph`` and ``pytorch_geometric.Data``. Supported conversion
80
+ formats can be retrieved from :const:`~graphein.ml.conversion.SUPPORTED_FORMATS`.
81
+
82
+ :param src_format: The type of graph you'd like to convert from. Supported formats are available in :const:`~graphein.ml.conversion.SUPPORTED_FORMATS`
83
+ :type src_format: Literal["nx", "pyg", "dgl", "jraph"]
84
+ :param dst_format: The type of graph format you'd like to convert to. Supported formats are available in:
85
+ ``graphein.ml.conversion.SUPPORTED_FORMATS``
86
+ :type dst_format: Literal["nx", "pyg", "dgl", "jraph"]
87
+ :param verbose: Select from ``"gnn"``, ``"default"``, ``"all_info"`` to determine how much information is preserved (features)
88
+ as some are unsupported by various downstream frameworks
89
+ :type verbose: graphein.ml.conversion.SUPPORTED_VERBOSITY
90
+ :param columns: List of columns in the node features to retain
91
+ :type columns: List[str], optional
92
+ """
93
+
94
+ def __init__(
95
+ self,
96
+ src_format: str,
97
+ dst_format: str,
98
+ verbose: SUPPORTED_VERBOSITY = "gnn",
99
+ columns: Optional[List[str]] = None,
100
+ ):
101
+ if (src_format not in SUPPORTED_FORMATS) or (
102
+ dst_format not in SUPPORTED_FORMATS
103
+ ):
104
+ raise ValueError(
105
+ "Please specify from supported format, "
106
+ + "/".join(SUPPORTED_FORMATS)
107
+ )
108
+ self.src_format = src_format
109
+ self.dst_format = dst_format
110
+
111
+ # supported_verbose_format = ["gnn", "default", "all_info"]
112
+ if (columns is None) and (verbose not in SUPPORTED_VERBOSITY):
113
+ raise ValueError(
114
+ "Please specify the supported verbose mode ("
115
+ + "/".join(SUPPORTED_VERBOSITY)
116
+ + ") or specify column names!"
117
+ )
118
+
119
+ if columns is None:
120
+ if verbose == "gnn":
121
+ columns = [
122
+ "edge_index",
123
+ "coords",
124
+ "dist_mat",
125
+ "name",
126
+ "node_id",
127
+ ]
128
+ elif verbose == "default":
129
+ columns = [
130
+ "b_factor",
131
+ "chain_id",
132
+ "coords",
133
+ "dist_mat",
134
+ "edge_index",
135
+ "kind",
136
+ "name",
137
+ "node_id",
138
+ "residue_name",
139
+ ]
140
+ elif verbose == "all_info":
141
+ columns = [
142
+ "atom_type",
143
+ "b_factor",
144
+ "chain_id",
145
+ "chain_ids",
146
+ "config",
147
+ "coords",
148
+ "dist_mat",
149
+ "edge_index",
150
+ "element_symbol",
151
+ "kind",
152
+ "name",
153
+ "node_id",
154
+ "node_type",
155
+ "pdb_df",
156
+ "raw_pdb_df",
157
+ "residue_name",
158
+ "residue_number",
159
+ "rgroup_df",
160
+ "sequence_A",
161
+ "sequence_B",
162
+ ]
163
+ self.columns = columns
164
+
165
+ self.type2form = {
166
+ "atom_type": "str",
167
+ "b_factor": "float",
168
+ "chain_id": "str",
169
+ "coords": "np.array",
170
+ "dist_mat": "np.array",
171
+ "element_symbol": "str",
172
+ "node_id": "str",
173
+ "residue_name": "str",
174
+ "residue_number": "int",
175
+ "edge_index": "torch.tensor",
176
+ "kind": "str",
177
+ }
178
+
179
+ def convert_nx_to_dgl(self, G: nx.Graph) -> dgl.DGLGraph:
180
+ """
181
+ Converts ``NetworkX`` graph to ``DGL``
182
+
183
+ :param G: ``nx.Graph`` to convert to ``DGLGraph``
184
+ :type G: nx.Graph
185
+ :return: ``DGLGraph`` object version of input ``NetworkX`` graph
186
+ :rtype: dgl.DGLGraph
187
+ """
188
+ g = dgl.DGLGraph()
189
+ node_id = list(G.nodes())
190
+ G = nx.convert_node_labels_to_integers(G)
191
+
192
+ ## add node level feat
193
+
194
+ node_dict = {}
195
+ for i, (_, feat_dict) in enumerate(G.nodes(data=True)):
196
+ for key, value in feat_dict.items():
197
+ if str(key) in self.columns:
198
+ node_dict[str(key)] = (
199
+ [value] if i == 0 else node_dict[str(key)] + [value]
200
+ )
201
+
202
+ string_dict = {}
203
+ node_dict_transformed = {}
204
+ for i, j in node_dict.items():
205
+ if i == "coords":
206
+ node_dict_transformed[i] = torch.Tensor(np.asarray(j)).type(
207
+ "torch.FloatTensor"
208
+ )
209
+ elif i == "dist_mat":
210
+ node_dict_transformed[i] = torch.Tensor(
211
+ np.asarray(j[0].values)
212
+ ).type("torch.FloatTensor")
213
+ elif self.type2form[i] == "str":
214
+ string_dict[i] = j
215
+ elif self.type2form[i] in ["float", "int"]:
216
+ node_dict_transformed[i] = torch.Tensor(np.array(j))
217
+ g.add_nodes(
218
+ len(node_id),
219
+ node_dict_transformed,
220
+ )
221
+
222
+ edge_dict = {}
223
+ edge_index = torch.LongTensor(list(G.edges)).t().contiguous()
224
+
225
+ # add edge level features
226
+ for i, (_, _, feat_dict) in enumerate(G.edges(data=True)):
227
+ for key, value in feat_dict.items():
228
+ if str(key) in self.columns:
229
+ edge_dict[str(key)] = (
230
+ list(value)
231
+ if i == 0
232
+ else edge_dict[str(key)] + list(value)
233
+ )
234
+
235
+ edge_transform_dict = {}
236
+ for i, j in node_dict.items():
237
+ if self.type2form[i] == "str":
238
+ string_dict[i] = j
239
+ elif self.type2form[i] in ["float", "int"]:
240
+ edge_transform_dict[i] = torch.Tensor(np.array(j))
241
+ g.add_edges(edge_index[0], edge_index[1], edge_transform_dict)
242
+
243
+ # add graph level features
244
+ graph_dict = {
245
+ str(feat_name): [G.graph[feat_name]]
246
+ for feat_name in G.graph
247
+ if str(feat_name) in self.columns
248
+ }
249
+
250
+ return g
251
+
252
+ def convert_nx_to_pyg(self, G: nx.Graph) -> Data:
253
+ """
254
+ Converts ``NetworkX`` graph to ``pytorch_geometric.data.Data`` object. Requires ``PyTorch Geometric`` (https://pytorch-geometric.readthedocs.io/en/latest/) to be installed.
255
+
256
+ :param G: ``nx.Graph`` to convert to PyTorch Geometric ``Data`` object
257
+ :type G: nx.Graph
258
+ :return: ``Data`` object containing networkx graph data
259
+ :rtype: pytorch_geometric.data.Data
260
+ """
261
+
262
+ # Initialise dict used to construct Data object & Assign node ids as a feature
263
+ data = {"node_id": list(G.nodes())}
264
+ G = nx.convert_node_labels_to_integers(G)
265
+
266
+ # Construct Edge Index
267
+ edge_index = torch.LongTensor(list(G.edges)).t().contiguous()
268
+
269
+ # Add node features
270
+ for i, (_, feat_dict) in enumerate(G.nodes(data=True)):
271
+ for key, value in feat_dict.items():
272
+ if str(key) in self.columns:
273
+ data[str(key)] = (
274
+ [value] if i == 0 else data[str(key)] + [value]
275
+ )
276
+
277
+ # Add edge features
278
+ for i, (_, _, feat_dict) in enumerate(G.edges(data=True)):
279
+ for key, value in feat_dict.items():
280
+ if str(key) in self.columns:
281
+ data[str(key)] = (
282
+ list(value) if i == 0 else data[str(key)] + list(value)
283
+ )
284
+
285
+ # Add graph-level features
286
+ for feat_name in G.graph:
287
+ if str(feat_name) in self.columns:
288
+ data[str(feat_name)] = [G.graph[feat_name]]
289
+
290
+ if "edge_index" in self.columns:
291
+ data["edge_index"] = edge_index.view(2, -1)
292
+
293
+ data = Data.from_dict(data)
294
+ data.num_nodes = G.number_of_nodes()
295
+ return data
296
+
297
+ @staticmethod
298
+ def convert_nx_to_nx(G: nx.Graph) -> nx.Graph:
299
+ """
300
+ Converts NetworkX graph (``nx.Graph``) to NetworkX graph (``nx.Graph``) object. Redundant - returns itself.
301
+
302
+ :param G: NetworkX Graph
303
+ :type G: nx.Graph
304
+ :return: NetworkX Graph
305
+ :rtype: nx.Graph
306
+ """
307
+ return G
308
+
309
+ @staticmethod
310
+ def convert_dgl_to_nx(G: dgl.DGLGraph) -> nx.Graph:
311
+ """
312
+ Converts a DGL Graph (``dgl.DGLGraph``) to a NetworkX (``nx.Graph``) object. Preserves node and edge attributes.
313
+
314
+ :param G: ``dgl.DGLGraph`` to convert to ``NetworkX`` graph.
315
+ :type G: dgl.DGLGraph
316
+ :return: NetworkX graph object.
317
+ :rtype: nx.Graph
318
+ """
319
+ node_attrs = G.node_attr_schemes().keys()
320
+ edge_attrs = G.edge_attr_schemes().keys()
321
+ return dgl.to_networkx(G, node_attrs, edge_attrs)
322
+
323
+ @staticmethod
324
+ def convert_pyg_to_nx(G: Data) -> nx.Graph:
325
+ """Converts PyTorch Geometric ``Data`` object to NetworkX graph (``nx.Graph``).
326
+
327
+ :param G: Pytorch Geometric Data.
328
+ :type G: torch_geometric.data.Data
329
+ :returns: NetworkX graph.
330
+ :rtype: nx.Graph
331
+ """
332
+ return torch_geometric.utils.to_networkx(G)
333
+
334
+ def convert_nx_to_jraph(self, G: nx.Graph) -> jraph.GraphsTuple:
335
+ """Converts NetworkX graph (``nx.Graph``) to Jraph GraphsTuple graph. Requires ``jax`` and ``Jraph``.
336
+
337
+ :param G: Networkx graph to convert.
338
+ :type G: nx.Graph
339
+ :return: Jraph GraphsTuple graph.
340
+ :rtype: jraph.GraphsTuple
341
+ """
342
+ G = nx.convert_node_labels_to_integers(G)
343
+
344
+ n_node = len(G)
345
+ n_edge = G.number_of_edges()
346
+ edge_list = list(G.edges())
347
+ senders, receivers = zip(*edge_list)
348
+ senders, receivers = jnp.array(senders), jnp.array(receivers)
349
+
350
+ # Add node features
351
+ node_features = {}
352
+ for i, (_, feat_dict) in enumerate(G.nodes(data=True)):
353
+ for key, value in feat_dict.items():
354
+ if str(key) in self.columns:
355
+ # node_features[str(key)] = (
356
+ # [value]
357
+ # if i == 0
358
+ # else node_features[str(key)] + [value]
359
+ # )
360
+ feat = (
361
+ [value]
362
+ if i == 0
363
+ else node_features[str(key)] + [value]
364
+ )
365
+ try:
366
+ feat = torch.tensor(feat)
367
+ node_features[str(key)] = feat
368
+ except TypeError:
369
+ node_features[str(key)] = feat
370
+
371
+ # Add edge features
372
+ edge_features = {}
373
+ for i, (_, _, feat_dict) in enumerate(G.edges(data=True)):
374
+ for key, value in feat_dict.items():
375
+ if str(key) in self.columns:
376
+ edge_features[str(key)] = (
377
+ list(value)
378
+ if i == 0
379
+ else edge_features[str(key)] + list(value)
380
+ )
381
+
382
+ # Add graph features
383
+ global_context = {
384
+ str(feat_name): [G.graph[feat_name]]
385
+ for feat_name in G.graph
386
+ if str(feat_name) in self.columns
387
+ }
388
+
389
+ return jraph.GraphsTuple(
390
+ nodes=node_features,
391
+ senders=senders,
392
+ receivers=receivers,
393
+ edges=edge_features,
394
+ n_node=n_node,
395
+ n_edge=n_edge,
396
+ globals=global_context,
397
+ )
398
+
399
+ def __call__(self, G: nx.Graph):
400
+ nx_g = eval("self.convert_" + self.src_format + "_to_nx(G)")
401
+ dst_g = eval("self.convert_nx_to_" + self.dst_format + "(nx_g)")
402
+ return dst_g
403
+
404
+
405
+ # def convert_nx_to_pyg_data(G: nx.Graph) -> Data:
406
+ # # Initialise dict used to construct Data object
407
+ # data = {"node_id": list(G.nodes())}
408
+
409
+ # G = nx.convert_node_labels_to_integers(G)
410
+
411
+ # # Construct Edge Index
412
+ # edge_index = torch.LongTensor(list(G.edges)).t().contiguous()
413
+
414
+ # # Add node features
415
+ # for i, (_, feat_dict) in enumerate(G.nodes(data=True)):
416
+ # for key, value in feat_dict.items():
417
+ # data[str(key)] = [value] if i == 0 else data[str(key)] + [value]
418
+
419
+ # # Add edge features
420
+ # for i, (_, _, feat_dict) in enumerate(G.edges(data=True)):
421
+ # for key, value in feat_dict.items():
422
+ # data[str(key)] = (
423
+ # list(value) if i == 0 else data[str(key)] + list(value)
424
+ # )
425
+
426
+ # # Add graph-level features
427
+ # for feat_name in G.graph:
428
+ # data[str(feat_name)] = [G.graph[feat_name]]
429
+
430
+ # data["edge_index"] = edge_index.view(2, -1)
431
+ # data = Data.from_dict(data)
432
+ # data.num_nodes = G.number_of_nodes()
433
+
434
+ # return data
435
+ def convert_nx_to_pyg_data(G: nx.Graph) -> Data:
436
+ # Initialise dict used to construct Data object
437
+ data = {"node_id": list(G.nodes())}
438
+
439
+ G = nx.convert_node_labels_to_integers(G)
440
+
441
+ # Construct Edge Index
442
+ edge_index = torch.LongTensor(list(G.edges)).t().contiguous()
443
+
444
+ # Add node features
445
+ for i, (_, feat_dict) in enumerate(G.nodes(data=True)):
446
+ for key, value in feat_dict.items():
447
+ data[str(key)] = [value] if i == 0 else data[str(key)] + [value]
448
+
449
+
450
+ # Add edge features
451
+ for i, (_, _, feat_dict) in enumerate(G.edges(data=True)):
452
+ for key, value in feat_dict.items():
453
+ if key == 'distance':
454
+ data[str(key)] = (
455
+ [value] if i == 0 else data[str(key)] + [value]
456
+ )
457
+ else:
458
+ data[str(key)] = (
459
+ [list(value)] if i == 0 else data[str(key)] + [list(value)]
460
+ )
461
+
462
+ # Add graph-level features
463
+ for feat_name in G.graph:
464
+ data[str(feat_name)] = [G.graph[feat_name]]
465
+
466
+ data["edge_index"] = edge_index.view(2, -1)
467
+ data = Data.from_dict(data)
468
+ data.num_nodes = G.number_of_nodes()
469
+
470
+ return data
graphs.py ADDED
@@ -0,0 +1,1137 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Functions for working with Protein Structure Graphs."""
2
+ # %%
3
+ # Graphein
4
+ # Author: Arian Jamasb <[email protected]>, Eric Ma, Charlie Harris
5
+ # License: MIT
6
+ # Project Website: https://github.com/a-r-j/graphein
7
+ # Code Repository: https://github.com/a-r-j/graphein
8
+ from __future__ import annotations
9
+
10
+ import logging
11
+ import traceback
12
+ from functools import partial
13
+ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
14
+
15
+ import networkx as nx
16
+ import numpy as np
17
+ import pandas as pd
18
+ # from Bio.PDB.Polypeptide import three_to_one
19
+ from biopandas.pdb import PandasPdb
20
+ from biopandas.mmcif import PandasMmcif
21
+ from rich.progress import Progress
22
+ from tqdm.contrib.concurrent import process_map
23
+
24
+ from graphein.protein.config import (
25
+ DSSPConfig,
26
+ GetContactsConfig,
27
+ ProteinGraphConfig,
28
+ )
29
+ from graphein.protein.edges.distance import (
30
+ add_distance_to_edges,
31
+ compute_distmat,
32
+ )
33
+ from graphein.protein.resi_atoms import BACKBONE_ATOMS, RESI_THREE_TO_1
34
+ from graphein.protein.subgraphs import extract_subgraph_from_chains
35
+ from graphein.protein.utils import (
36
+ ProteinGraphConfigurationError,
37
+ compute_rgroup_dataframe,
38
+ filter_dataframe,
39
+ get_protein_name_from_filename,
40
+ three_to_one_with_mods,
41
+ )
42
+ from graphein.rna.constants import RNA_ATOMS
43
+ from graphein.utils.utils import (
44
+ annotate_edge_metadata,
45
+ annotate_graph_metadata,
46
+ annotate_node_metadata,
47
+ compute_edges,
48
+ )
49
+
50
+ from .utils_convert import biopandas_mmcif2pdb
51
+
52
+ # logging.basicConfig(level="DEBUG")
53
+ log = logging.getLogger(__name__)
54
+
55
+
56
+
57
+ def subset_structure_to_rna(
58
+ df: pd.DataFrame,
59
+ ) -> pd.DataFrame:
60
+ """
61
+ Return a subset of atomic dataframe that contains only certain atom names relevant for RNA structures.
62
+
63
+ :param df: Protein Structure dataframe to subset
64
+ :type df: pd.DataFrame
65
+ :returns: Subsetted protein structure dataframe
66
+ :rtype: pd.DataFrame
67
+ """
68
+ return filter_dataframe(
69
+ df, by_column="atom_name", list_of_values=RNA_ATOMS, boolean=True
70
+ )
71
+
72
+
73
+ def read_pdb_to_dataframe(
74
+ pdb_path: Optional[str] = None,
75
+ pdb_code: Optional[str] = None,
76
+ uniprot_id: Optional[str] = None,
77
+ model_index: int = 1,
78
+ ) -> pd.DataFrame:
79
+ """
80
+ Reads PDB file to ``PandasPDB`` object.
81
+
82
+ Returns ``atomic_df``, which is a dataframe enumerating all atoms and their cartesian coordinates in 3D space. Also
83
+ contains associated metadata from the PDB file.
84
+
85
+ :param pdb_path: path to PDB file. Defaults to ``None``.
86
+ :type pdb_path: str, optional
87
+ :param pdb_code: 4-character PDB accession. Defaults to ``None``.
88
+ :type pdb_code: str, optional
89
+ :param uniprot_id: UniProt ID to build graph from AlphaFoldDB. Defaults to ``None``.
90
+ :type uniprot_id: str, optional
91
+ :param model_index: Index of model to read. Only relevant for structures containing ensembles. Defaults to ``1``.
92
+ :type model_index: int, optional
93
+ :param verbose: print dataframe?
94
+ :type verbose: bool
95
+ :param granularity: Specifies granularity of dataframe. See :class:`~graphein.protein.config.ProteinGraphConfig` for further
96
+ details.
97
+ :type granularity: str
98
+ :returns: ``pd.DataFrame`` containing protein structure
99
+ :rtype: pd.DataFrame
100
+ """
101
+ if pdb_code is None and pdb_path is None and uniprot_id is None:
102
+ raise NameError(
103
+ "One of pdb_code, pdb_path or uniprot_id must be specified!"
104
+ )
105
+
106
+ if pdb_path is not None:
107
+ if pdb_path.endswith('cif'):
108
+ atomic_df = PandasMmcif().read_mmcif(pdb_path)
109
+ atomic_df = biopandas_mmcif2pdb(atomic_df, model_index)
110
+ else:
111
+ atomic_df = PandasPdb().read_pdb(pdb_path)
112
+ else:
113
+ if uniprot_id is not None:
114
+ atomic_df = PandasPdb().fetch_pdb(
115
+ uniprot_id=uniprot_id, source="alphafold2-v2"
116
+ )
117
+ else:
118
+ atomic_df = PandasPdb().fetch_pdb(pdb_code)
119
+
120
+ atomic_df = atomic_df.get_model(model_index)
121
+ if len(atomic_df.df["ATOM"]) == 0:
122
+ raise ValueError(f"No model found for index: {model_index}")
123
+
124
+ return pd.concat([atomic_df.df["ATOM"], atomic_df.df["HETATM"]])
125
+
126
+
127
+ def label_node_id(df: pd.DataFrame, granularity: str) -> pd.DataFrame:
128
+ df["node_id"] = (
129
+ df["chain_id"].apply(str)
130
+ + ":"
131
+ + df["residue_name"]
132
+ + ":"
133
+ + df["residue_number"].apply(str)
134
+ )
135
+ df["residue_id"] = df["node_id"]
136
+ if granularity == "atom":
137
+ df["node_id"] = df["node_id"] + ":" + df["atom_name"]
138
+ elif granularity in {"rna_atom", "rna_centroid"}:
139
+ df["node_id"] = (
140
+ df["node_id"]
141
+ + ":"
142
+ + df["atom_number"].apply(str)
143
+ + ":"
144
+ + df["atom_name"]
145
+ )
146
+ return df
147
+
148
+
149
+ def deprotonate_structure(df: pd.DataFrame) -> pd.DataFrame:
150
+ """Remove protons from PDB dataframe.
151
+
152
+ :param df: Atomic dataframe.
153
+ :type df: pd.DataFrame
154
+ :returns: Atomic dataframe with all ``atom_name == "H"`` removed.
155
+ :rtype: pd.DataFrame
156
+ """
157
+ log.debug(
158
+ "Deprotonating protein. This removes H atoms from the pdb_df dataframe"
159
+ )
160
+ return filter_dataframe(
161
+ df, by_column="element_symbol", list_of_values=["H"], boolean=False
162
+ )
163
+
164
+
165
+ def convert_structure_to_centroids(df: pd.DataFrame) -> pd.DataFrame:
166
+ """Overwrite existing ``(x, y, z)`` coordinates with centroids of the amino acids.
167
+
168
+ :param df: Pandas Dataframe protein structure to convert into a dataframe of centroid positions.
169
+ :type df: pd.DataFrame
170
+ :return: pd.DataFrame with atoms/residues positions converted into centroid positions.
171
+ :rtype: pd.DataFrame
172
+ """
173
+ log.debug(
174
+ "Converting dataframe to centroids. This averages XYZ coords of the atoms in a residue"
175
+ )
176
+
177
+ centroids = calculate_centroid_positions(df)
178
+ df = df.loc[df["atom_name"] == "CA"].reset_index(drop=True)
179
+ df["x_coord"] = centroids["x_coord"]
180
+ df["y_coord"] = centroids["y_coord"]
181
+ df["z_coord"] = centroids["z_coord"]
182
+
183
+ return df
184
+
185
+
186
+ def subset_structure_to_atom_type(
187
+ df: pd.DataFrame, granularity: str
188
+ ) -> pd.DataFrame:
189
+ """
190
+ Return a subset of atomic dataframe that contains only certain atom names.
191
+
192
+ :param df: Protein Structure dataframe to subset.
193
+ :type df: pd.DataFrame
194
+ :returns: Subsetted protein structure dataframe.
195
+ :rtype: pd.DataFrame
196
+ """
197
+ return filter_dataframe(
198
+ df, by_column="atom_name", list_of_values=[granularity], boolean=True
199
+ )
200
+
201
+
202
+ def remove_insertions(df: pd.DataFrame, keep: str = "first") -> pd.DataFrame:
203
+ """
204
+ This function removes insertions from PDB dataframes.
205
+
206
+ :param df: Protein Structure dataframe to remove insertions from.
207
+ :type df: pd.DataFrame
208
+ :param keep: Specifies which insertion to keep. Options are ``"first"`` or ``"last"``.
209
+ Default is ``"first"``
210
+ :type keep: str
211
+ :return: Protein structure dataframe with insertions removed
212
+ :rtype: pd.DataFrame
213
+ """
214
+ # Catches unnamed insertions
215
+ duplicates = df.duplicated(
216
+ subset=["chain_id", "residue_number", "atom_name"], keep=keep
217
+ )
218
+ df = df[~duplicates]
219
+
220
+ # Catches explicit insertions
221
+ df = filter_dataframe(
222
+ df, by_column="insertion", list_of_values=[""], boolean=True
223
+ )
224
+
225
+ # Remove alt_locs
226
+ df = filter_dataframe(
227
+ df, by_column="alt_loc", list_of_values=["", "A"], boolean=True
228
+ )
229
+
230
+ return df
231
+
232
+
233
+ def filter_hetatms(
234
+ df: pd.DataFrame, keep_hets: List[str]
235
+ ) -> List[pd.DataFrame]:
236
+ """Return hetatms of interest.
237
+
238
+ :param df: Protein Structure dataframe to filter hetatoms from.
239
+ :type df: pd.DataFrame
240
+ :param keep_hets: List of hetero atom names to keep.
241
+ :returns: Protein structure dataframe with heteroatoms removed
242
+ :rtype: pd.DataFrame
243
+ """
244
+ return [df.loc[df["residue_name"] == hetatm] for hetatm in keep_hets]
245
+
246
+
247
+ def process_dataframe(
248
+ protein_df: pd.DataFrame,
249
+ atom_df_processing_funcs: Optional[List[Callable]] = None,
250
+ hetatom_df_processing_funcs: Optional[List[Callable]] = None,
251
+ granularity: str = "centroids",
252
+ chain_selection: str = "all",
253
+ insertions: bool = False,
254
+ deprotonate: bool = True,
255
+ keep_hets: List[str] = [],
256
+ verbose: bool = False,
257
+ ) -> pd.DataFrame:
258
+ """
259
+ Process ATOM and HETATM dataframes to produce singular dataframe used for graph construction.
260
+
261
+ :param protein_df: Dataframe to process.
262
+ Should be the object returned from :func:`~graphein.protein.graphs.read_pdb_to_dataframe`.
263
+ :type protein_df: pd.DataFrame
264
+ :param atom_df_processing_funcs: List of functions to process dataframe. These must take in a dataframe and return a
265
+ dataframe. Defaults to None.
266
+ :type atom_df_processing_funcs: List[Callable], optional
267
+ :param hetatom_df_processing_funcs: List of functions to process the hetatom dataframe. These must take in a dataframe and return a dataframe
268
+ :type hetatom_df_processing_funcs: List[Callable], optional
269
+ :param granularity: The level of granularity for the graph. This determines the node definition.
270
+ Acceptable values include: ``"centroids"``, ``"atoms"``,
271
+ any of the atom_names in the PDB file (e.g. ``"CA"``, ``"CB"``, ``"OG"``, etc.).
272
+ See: :const:`~graphein.protein.config.GRAPH_ATOMS` and :const:`~graphein.protein.config.GRANULARITY_OPTS`.
273
+ :type granularity: str
274
+ :param insertions: Whether or not to keep insertions.
275
+ :param insertions: bool
276
+ :param deprotonate: Whether or not to remove hydrogen atoms (i.e. deprotonation).
277
+ :type deprotonate: bool
278
+ :param keep_hets: Hetatoms to keep. Defaults to an empty list.
279
+ To keep a hetatom, pass it inside a list of hetatom names to keep.
280
+ :type keep_hets: List[str]
281
+ :param verbose: Verbosity level.
282
+ :type verbose: bool
283
+ :param chain_selection: Which protein chain to select. Defaults to ``"all"``. Eg can use ``"ACF"``
284
+ to select 3 chains (``A``, ``C`` & ``F``)
285
+ :type chain_selection: str
286
+ :return: A protein dataframe that can be consumed by
287
+ other graph construction functions.
288
+ :rtype: pd.DataFrame
289
+ """
290
+ protein_df = label_node_id(protein_df, granularity=granularity)
291
+ # TODO: Need to properly define what "granularity" is supposed to do.
292
+ atoms = filter_dataframe(
293
+ protein_df,
294
+ by_column="record_name",
295
+ list_of_values=["ATOM"],
296
+ boolean=True,
297
+ )
298
+ hetatms = filter_dataframe(
299
+ protein_df,
300
+ by_column="record_name",
301
+ list_of_values=["HETATM"],
302
+ boolean=True,
303
+ )
304
+
305
+ # This block enables processing via a list of supplied functions operating on the atom and hetatom dataframes
306
+ # If these are provided, the dataframe returned will be computed only from these and the default workflow
307
+ # below this block will not execute.
308
+ if atom_df_processing_funcs is not None:
309
+ for func in atom_df_processing_funcs:
310
+ atoms = func(atoms)
311
+ if hetatom_df_processing_funcs is None:
312
+ return atoms
313
+
314
+ if hetatom_df_processing_funcs is not None:
315
+ for func in hetatom_df_processing_funcs:
316
+ hetatms = func(hetatms)
317
+ return pd.concat([atoms, hetatms])
318
+
319
+ if keep_hets:
320
+ hetatms_to_keep = filter_hetatms(hetatms, keep_hets)
321
+ atoms = pd.concat([atoms] + hetatms_to_keep)
322
+
323
+ # Deprotonate structure by removing H atoms
324
+ if deprotonate:
325
+ atoms = deprotonate_structure(atoms)
326
+
327
+ # Restrict DF to desired granularity
328
+ if granularity == "atom":
329
+ pass
330
+ elif granularity in {"centroids", "rna_centroid"}:
331
+ atoms = convert_structure_to_centroids(atoms)
332
+ elif granularity == "rna_atom":
333
+ atoms = subset_structure_to_rna(atoms)
334
+ else:
335
+ atoms = subset_structure_to_atom_type(atoms, granularity)
336
+
337
+ protein_df = atoms
338
+
339
+ # Remove alt_loc residues
340
+ if not insertions:
341
+ protein_df = remove_insertions(protein_df)
342
+
343
+ # perform chain selection
344
+ protein_df = select_chains(
345
+ protein_df, chain_selection=chain_selection, verbose=verbose
346
+ )
347
+
348
+ log.debug(f"Detected {len(protein_df)} total nodes")
349
+
350
+ # Sort dataframe to place HETATMs
351
+ protein_df = sort_dataframe(protein_df)
352
+
353
+ return protein_df
354
+
355
+
356
+ def sort_dataframe(df: pd.DataFrame) -> pd.DataFrame:
357
+ """Sorts a protein dataframe by chain->residue number->atom number
358
+
359
+ This is useful for distributing hetatms/modified residues through the DF.
360
+
361
+ :param df: Protein dataframe to sort.
362
+ :type df: pd.DataFrame
363
+ :return: Sorted protein dataframe.
364
+ :rtype: pd.DataFrame
365
+ """
366
+ return df.sort_values(by=["chain_id", "residue_number", "atom_number"])
367
+
368
+
369
+ def assign_node_id_to_dataframe(
370
+ protein_df: pd.DataFrame, granularity: str
371
+ ) -> pd.DataFrame:
372
+ """
373
+ Assigns the node ID back to the ``pdb_df`` dataframe
374
+
375
+ :param protein_df: Structure Dataframe
376
+ :type protein_df: pd.DataFrame
377
+ :param granularity: Granularity of graph. Atom-level,
378
+ residue (e.g. ``CA``) or ``centroids``.
379
+ See: :const:`~graphein.protein.config.GRAPH_ATOMS`
380
+ and :const:`~graphein.protein.config.GRANULARITY_OPTS`.
381
+ :type granularity: str
382
+ :return: Returns dataframe with added ``node_ids``
383
+ :rtype: pd.DataFrame
384
+ """
385
+ protein_df["node_id"] = (
386
+ protein_df["chain_id"].apply(str)
387
+ + ":"
388
+ + protein_df["residue_name"]
389
+ + ":"
390
+ + protein_df["residue_number"].apply(str)
391
+ )
392
+ if granularity in {"atom", "rna_atom"}:
393
+ protein_df[
394
+ "node_id"
395
+ ] = f'{protein_df["node_id"]}:{protein_df["atom_name"]}'
396
+
397
+
398
+ def select_chains(
399
+ protein_df: pd.DataFrame, chain_selection: str, verbose: bool = False
400
+ ) -> pd.DataFrame:
401
+ """
402
+ Extracts relevant chains from ``protein_df``.
403
+
404
+ :param protein_df: pandas dataframe of PDB subsetted to relevant atoms
405
+ (``CA``, ``CB``).
406
+ :type protein_df: pd.DataFrame
407
+ :param chain_selection: Specifies chains that should be extracted from
408
+ the larger complexed structure.
409
+ :type chain_selection: str
410
+ :param verbose: Print dataframe?
411
+ :type verbose: bool
412
+ :return: Protein structure dataframe containing only entries in the
413
+ chain selection.
414
+ :rtype: pd.DataFrame
415
+ """
416
+ if chain_selection != "all":
417
+ protein_df = filter_dataframe(
418
+ protein_df,
419
+ by_column="chain_id",
420
+ list_of_values=list(chain_selection),
421
+ boolean=True,
422
+ )
423
+
424
+ return protein_df
425
+
426
+
427
+ def initialise_graph_with_metadata(
428
+ protein_df: pd.DataFrame,
429
+ raw_pdb_df: pd.DataFrame,
430
+ granularity: str,
431
+ name: Optional[str] = None,
432
+ pdb_code: Optional[str] = None,
433
+ pdb_path: Optional[str] = None,
434
+ ) -> nx.Graph:
435
+ """
436
+ Initializes the nx Graph object with initial metadata.
437
+
438
+ :param protein_df: Processed Dataframe of protein structure.
439
+ :type protein_df: pd.DataFrame
440
+ :param raw_pdb_df: Unprocessed dataframe of protein structure for comparison and traceability downstream.
441
+ :type raw_pdb_df: pd.DataFrame
442
+ :param granularity: Granularity of the graph (eg ``"atom"``, ``"CA"``, ``"CB"`` etc or ``"centroid"``).
443
+ See: :const:`~graphein.protein.config.GRAPH_ATOMS` and :const:`~graphein.protein.config.GRANULARITY_OPTS`.
444
+ :type granularity: str
445
+ :param name: specified given name for the graph. If None, the PDB code or the file name will be used to name the graph.
446
+ :type name: Optional[str], defaults to ``None``
447
+ :param pdb_code: PDB ID / Accession code, if the PDB is available on the PDB database.
448
+ :type pdb_code: Optional[str], defaults to ``None``
449
+ :param pdb_path: path to local PDB file, if constructing a graph from a local file.
450
+ :type pdb_path: Optional[str], defaults to ``None``
451
+ :return: Returns initial protein structure graph with metadata.
452
+ :rtype: nx.Graph
453
+ """
454
+
455
+ # Get name for graph if no name was provided
456
+ if name is None:
457
+ if pdb_path is not None:
458
+ name = get_protein_name_from_filename(pdb_path)
459
+ else:
460
+ name = pdb_code
461
+
462
+ G = nx.Graph(
463
+ name=name,
464
+ pdb_code=pdb_code,
465
+ pdb_path=pdb_path,
466
+ chain_ids=list(protein_df["chain_id"].unique()),
467
+ pdb_df=protein_df,
468
+ raw_pdb_df=raw_pdb_df,
469
+ rgroup_df=compute_rgroup_dataframe(remove_insertions(raw_pdb_df)),
470
+ coords=np.asarray(protein_df[["x_coord", "y_coord", "z_coord"]]),
471
+ )
472
+
473
+ # Create graph and assign intrinsic graph-level metadata
474
+ G.graph["node_type"] = granularity
475
+
476
+ # Add Sequences to graph metadata
477
+ for c in G.graph["chain_ids"]:
478
+ if granularity == "rna_atom":
479
+ sequence = protein_df.loc[protein_df["chain_id"] == c][
480
+ "residue_name"
481
+ ].str.cat()
482
+ else:
483
+ sequence = (
484
+ protein_df.loc[protein_df["chain_id"] == c]["residue_name"]
485
+ .apply(three_to_one_with_mods)
486
+ .str.cat()
487
+ )
488
+ G.graph[f"sequence_{c}"] = sequence
489
+ return G
490
+
491
+
492
+ def add_nodes_to_graph(
493
+ G: nx.Graph,
494
+ protein_df: Optional[pd.DataFrame] = None,
495
+ verbose: bool = False,
496
+ ) -> nx.Graph:
497
+ """Add nodes into protein graph.
498
+
499
+ :param G: ``nx.Graph`` with metadata to populate with nodes.
500
+ :type G: nx.Graph
501
+ :protein_df: DataFrame of protein structure containing nodes & initial node metadata to add to the graph.
502
+ :type protein_df: pd.DataFrame, optional
503
+ :param verbose: Controls verbosity of this step.
504
+ :type verbose: bool
505
+ :returns: nx.Graph with nodes added.
506
+ :rtype: nx.Graph
507
+ """
508
+
509
+ # If no protein dataframe is supplied, use the one stored in the Graph object
510
+ if protein_df is None:
511
+ protein_df = G.graph["pdb_df"]
512
+ # Assign intrinsic node attributes
513
+ chain_id = protein_df["chain_id"].apply(str)
514
+ residue_name = protein_df["residue_name"]
515
+ residue_number = protein_df["residue_number"] # .apply(str)
516
+ coords = np.asarray(protein_df[["x_coord", "y_coord", "z_coord"]])
517
+ b_factor = protein_df["b_factor"]
518
+ atom_type = protein_df["atom_name"]
519
+ nodes = protein_df["node_id"]
520
+ element_symbol = protein_df["element_symbol"]
521
+ G.add_nodes_from(nodes)
522
+
523
+ # Set intrinsic node attributes
524
+ nx.set_node_attributes(G, dict(zip(nodes, chain_id)), "chain_id")
525
+ nx.set_node_attributes(G, dict(zip(nodes, residue_name)), "residue_name")
526
+ nx.set_node_attributes(
527
+ G, dict(zip(nodes, residue_number)), "residue_number"
528
+ )
529
+ nx.set_node_attributes(G, dict(zip(nodes, atom_type)), "atom_type")
530
+ nx.set_node_attributes(
531
+ G, dict(zip(nodes, element_symbol)), "element_symbol"
532
+ )
533
+ nx.set_node_attributes(G, dict(zip(nodes, coords)), "coords")
534
+ nx.set_node_attributes(G, dict(zip(nodes, b_factor)), "b_factor")
535
+
536
+ # TODO: include charge, line_idx for traceability?
537
+ if verbose:
538
+ print(nx.info(G))
539
+ print(G.nodes())
540
+
541
+ return G
542
+
543
+
544
+ def calculate_centroid_positions(
545
+ atoms: pd.DataFrame, verbose: bool = False
546
+ ) -> pd.DataFrame:
547
+ """
548
+ Calculates position of sidechain centroids.
549
+
550
+ :param atoms: ATOM df of protein structure.
551
+ :type atoms: pd.DataFrame
552
+ :param verbose: bool controlling verbosity.
553
+ :type verbose: bool
554
+ :return: centroids (df).
555
+ :rtype: pd.DataFrame
556
+ """
557
+ centroids = (
558
+ atoms.groupby("residue_number")
559
+ .mean()[["x_coord", "y_coord", "z_coord"]]
560
+ .reset_index()
561
+ )
562
+ if verbose:
563
+ print(f"Calculated {len(centroids)} centroid nodes")
564
+ log.debug(f"Calculated {len(centroids)} centroid nodes")
565
+ return centroids
566
+
567
+
568
+ def compute_edges(
569
+ G: nx.Graph,
570
+ funcs: List[Callable],
571
+ get_contacts_config: Optional[GetContactsConfig] = None,
572
+ ) -> nx.Graph:
573
+ """
574
+ Computes edges for the protein structure graph. Will compute a pairwise
575
+ distance matrix between nodes which is
576
+ added to the graph metadata to facilitate some edge computations.
577
+
578
+ :param G: nx.Graph with nodes to add edges to.
579
+ :type G: nx.Graph
580
+ :param funcs: List of edge construction functions.
581
+ :type funcs: List[Callable]
582
+ :param get_contacts_config: Config object for ``GetContacts`` if
583
+ intramolecular edges are being used.
584
+ :type get_contacts_config: graphein.protein.config.GetContactsConfig
585
+ :return: Graph with added edges.
586
+ :rtype: nx.Graph
587
+ """
588
+ # This control flow prevents unnecessary computation of the distance matrices
589
+ if "config" in G.graph:
590
+ if G.graph["config"].granularity == "atom":
591
+ G.graph["atomic_dist_mat"] = compute_distmat(G.graph["pdb_df"])
592
+ else:
593
+ G.graph["dist_mat"] = compute_distmat(G.graph["pdb_df"])
594
+
595
+ for func in funcs:
596
+ func(G)
597
+
598
+ return add_distance_to_edges(G)
599
+
600
+
601
+ def construct_graph(
602
+ config: Optional[ProteinGraphConfig] = None,
603
+ name: Optional[str] = None,
604
+ pdb_path: Optional[str] = None,
605
+ uniprot_id: Optional[str] = None,
606
+ pdb_code: Optional[str] = None,
607
+ chain_selection: str = "all",
608
+ model_index: int = 1,
609
+ df_processing_funcs: Optional[List[Callable]] = None,
610
+ edge_construction_funcs: Optional[List[Callable]] = None,
611
+ edge_annotation_funcs: Optional[List[Callable]] = None,
612
+ node_annotation_funcs: Optional[List[Callable]] = None,
613
+ graph_annotation_funcs: Optional[List[Callable]] = None,
614
+ ) -> nx.Graph:
615
+ """
616
+ Constructs protein structure graph from a ``pdb_code`` or ``pdb_path``.
617
+
618
+ Users can provide a :class:`~graphein.protein.config.ProteinGraphConfig`
619
+ object to specify construction parameters.
620
+
621
+ However, config parameters can be overridden by passing arguments directly to the function.
622
+
623
+ :param config: :class:`~graphein.protein.config.ProteinGraphConfig` object. If None, defaults to config in ``graphein.protein.config``.
624
+ :type config: graphein.protein.config.ProteinGraphConfig, optional
625
+ :param name: an optional given name for the graph. the PDB ID or PDB file name will be used if not specified.
626
+ :type name: str, optional
627
+ :param pdb_path: Path to ``pdb_file`` when constructing a graph from a local pdb file. Default is ``None``.
628
+ :type pdb_path: Optional[str], defaults to ``None``
629
+ :param pdb_code: A 4-character PDB ID / accession to be used to construct the graph, if available. Default is ``None``.
630
+ :type pdb_code: Optional[str], defaults to ``None``
631
+ :param uniprot_id: UniProt accession ID to build graph from AlphaFold2DB. Default is ``None``.
632
+ :type uniprot_id: str, optional
633
+ :param chain_selection: String of polypeptide chains to include in graph. E.g ``"ABDF"`` or ``"all"``. Default is ``"all"``.
634
+ :type chain_selection: str
635
+ :param model_index: Index of model to use in the case of structural ensembles. Default is ``1``.
636
+ :type model_index: int
637
+ :param df_processing_funcs: List of dataframe processing functions. Default is ``None``.
638
+ :type df_processing_funcs: List[Callable], optional
639
+ :param edge_construction_funcs: List of edge construction functions. Default is ``None``.
640
+ :type edge_construction_funcs: List[Callable], optional
641
+ :param edge_annotation_funcs: List of edge annotation functions. Default is ``None``.
642
+ :type edge_annotation_funcs: List[Callable], optional
643
+ :param node_annotation_funcs: List of node annotation functions. Default is ``None``.
644
+ :type node_annotation_funcs: List[Callable], optional
645
+ :param graph_annotation_funcs: List of graph annotation function. Default is ``None``.
646
+ :type graph_annotation_funcs: List[Callable]
647
+ :return: Protein Structure Graph
648
+ :rtype: nx.Graph
649
+ """
650
+
651
+ if pdb_code is None and pdb_path is None and uniprot_id is None:
652
+ raise ValueError(
653
+ "Either a PDB ID, UniProt ID or a path to a local PDB file"
654
+ " must be specified to construct a graph"
655
+ )
656
+
657
+ # If no config is provided, use default
658
+ if config is None:
659
+ config = ProteinGraphConfig()
660
+ with Progress(transient=True) as progress:
661
+ task1 = progress.add_task("Reading PDB file...", total=1)
662
+ # Get name from pdb_file is no pdb_code is provided
663
+ # if pdb_path and (pdb_code is None and uniprot_id is None):
664
+ # pdb_code = get_protein_name_from_filename(pdb_path)
665
+ # pdb_code = pdb_code if len(pdb_code) == 4 else None
666
+ progress.advance(task1)
667
+
668
+ # If config params are provided, overwrite them
669
+ config.protein_df_processing_functions = (
670
+ df_processing_funcs
671
+ if config.protein_df_processing_functions is None
672
+ else config.protein_df_processing_functions
673
+ )
674
+ config.edge_construction_functions = (
675
+ edge_construction_funcs
676
+ if config.edge_construction_functions is None
677
+ else config.edge_construction_functions
678
+ )
679
+ config.node_metadata_functions = (
680
+ node_annotation_funcs
681
+ if config.node_metadata_functions is None
682
+ else config.node_metadata_functions
683
+ )
684
+ config.graph_metadata_functions = (
685
+ graph_annotation_funcs
686
+ if config.graph_metadata_functions is None
687
+ else config.graph_metadata_functions
688
+ )
689
+ config.edge_metadata_functions = (
690
+ edge_annotation_funcs
691
+ if config.edge_metadata_functions is None
692
+ else config.edge_metadata_functions
693
+ )
694
+
695
+ raw_df = read_pdb_to_dataframe(
696
+ pdb_path,
697
+ pdb_code,
698
+ uniprot_id,
699
+ model_index=model_index,
700
+ )
701
+
702
+
703
+ task2 = progress.add_task("Processing PDB dataframe...", total=1)
704
+ # raw_df = label_node_id(raw_df, granularity=config.granularity)
705
+ # raw_df.df["ATOM"] = label_node_id(
706
+ # raw_df.df["ATOM"], granularity=config.granularity
707
+ # )
708
+ # raw_df.df["HETATM"] = label_node_id(
709
+ # raw_df.df["HETATM"], granularity=config.granularity
710
+ # )
711
+ raw_df = sort_dataframe(raw_df)
712
+ protein_df = process_dataframe(
713
+ raw_df,
714
+ chain_selection=chain_selection,
715
+ granularity=config.granularity,
716
+ insertions=config.insertions,
717
+ keep_hets=config.keep_hets,
718
+ )
719
+ progress.advance(task2)
720
+
721
+ task3 = progress.add_task("Initializing graph...", total=1)
722
+ # Initialise graph with metadata
723
+ g = initialise_graph_with_metadata(
724
+ protein_df=protein_df,
725
+ raw_pdb_df=raw_df,
726
+ name=name,
727
+ pdb_code=pdb_code,
728
+ pdb_path=pdb_path,
729
+ granularity=config.granularity,
730
+ )
731
+ # Add nodes to graph
732
+ g = add_nodes_to_graph(g)
733
+ # Add config to graph
734
+ g.graph["config"] = config
735
+ g.graph["path"] = g.graph["pdb_path"]
736
+
737
+ # Annotate additional node metadata
738
+ if config.node_metadata_functions is not None:
739
+ g = annotate_node_metadata(g, config.node_metadata_functions)
740
+ progress.advance(task3)
741
+ task4 = progress.add_task("Constructing edges...", total=1)
742
+ # Compute graph edges
743
+ g = compute_edges(
744
+ g,
745
+ funcs=config.edge_construction_functions,
746
+ get_contacts_config=None,
747
+ )
748
+ progress.advance(task4)
749
+
750
+ # Annotate additional graph metadata
751
+ # print(g.graph['dssp_df'])
752
+ if config.graph_metadata_functions is not None:
753
+ g = annotate_graph_metadata(g, config.graph_metadata_functions)
754
+
755
+ # Annotate additional edge metadata
756
+ if config.edge_metadata_functions is not None:
757
+ g = annotate_edge_metadata(g, config.edge_metadata_functions)
758
+
759
+ return g
760
+
761
+
762
+ def _mp_graph_constructor(
763
+ args: Tuple[str, str, int], source: str, config: ProteinGraphConfig
764
+ ) -> Union[nx.Graph, None]:
765
+ """
766
+ Protein graph constructor for use in multiprocessing several protein structure graphs.
767
+
768
+ :param args: Tuple of pdb code/path and the chain selection for that PDB.
769
+ :type args: Tuple[str, str]
770
+ :param use_pdb_code: Whether we are using ``"pdb_code"``s, ``pdb_path``s or ``"uniprot_id"``s.
771
+ :type use_pdb_code: bool
772
+ :param config: Protein structure graph construction config (see: :class:`graphein.protein.config.ProteinGraphConfig`).
773
+ :type config: ProteinGraphConfig
774
+ :return: Protein structure graph or ``None`` if an error is encountered.
775
+ :rtype: Union[nx.Graph, None]
776
+ """
777
+ log.info(
778
+ f"Constructing graph for: {args[0]}. Chain selection: {args[1]}. Model index: {args[2]}"
779
+ )
780
+ func = partial(construct_graph, config=config)
781
+ try:
782
+ if source == "pdb_code":
783
+ return func(
784
+ pdb_code=args[0], chain_selection=args[1], model_index=args[2]
785
+ )
786
+ elif source == "pdb_path":
787
+ return func(
788
+ pdb_path=args[0], chain_selection=args[1], model_index=args[2]
789
+ )
790
+ elif source == "uniprot_id":
791
+ return func(
792
+ uniprot_id=args[0],
793
+ chain_selection=args[1],
794
+ model_index=args[2],
795
+ )
796
+
797
+ except Exception as ex:
798
+ log.info(
799
+ f"Graph construction error (PDB={args[0]})! {traceback.format_exc()}"
800
+ )
801
+ log.info(ex)
802
+ return None
803
+
804
+
805
+ def construct_graphs_mp(
806
+ pdb_code_it: Optional[List[str]] = None,
807
+ pdb_path_it: Optional[List[str]] = None,
808
+ uniprot_id_it: Optional[List[str]] = None,
809
+ chain_selections: Optional[List[str]] = None,
810
+ model_indices: Optional[List[str]] = None,
811
+ config: ProteinGraphConfig = ProteinGraphConfig(),
812
+ num_cores: int = 16,
813
+ return_dict: bool = True,
814
+ out_path: Optional[str] = None,
815
+ ) -> Union[List[nx.Graph], Dict[str, nx.Graph]]:
816
+ """
817
+ Constructs protein graphs for a list of pdb codes or pdb paths using multiprocessing.
818
+
819
+ :param pdb_code_it: List of pdb codes to use for protein graph construction
820
+ :type pdb_code_it: Optional[List[str]], defaults to ``None``
821
+ :param pdb_path_it: List of paths to PDB files to use for protein graph construction
822
+ :type pdb_path_it: Optional[List[str]], defaults to ``None``
823
+ :param chain_selections: List of chains to select from the protein structures (e.g. ``["ABC", "A", "L", "CD"...]``)
824
+ :type chain_selections: Optional[List[str]], defaults to ``None``
825
+ :param model_indices: List of model indices to use for protein graph construction. Only relevant for structures containing ensembles of models.
826
+ :type model_indices: Optional[List[str]], defaults to ``None``
827
+ :param config: ProteinGraphConfig to use.
828
+ :type config: graphein.protein.config.ProteinGraphConfig, defaults to default config params
829
+ :param num_cores: Number of cores to use for multiprocessing. The more the merrier
830
+ :type num_cores: int, defaults to ``16``
831
+ :param return_dict: Whether or not to return a dictionary (indexed by pdb codes/paths) or a list of graphs.
832
+ :type return_dict: bool, default to ``True``
833
+ :param out_path: Path to save the graphs to. If None, graphs are not saved.
834
+ :type out_path: Optional[str], defaults to ``None``
835
+ :return: Iterable of protein graphs. None values indicate there was a problem in constructing the graph for this particular pdb
836
+ :rtype: Union[List[nx.Graph], Dict[str, nx.Graph]]
837
+ """
838
+ assert (
839
+ pdb_code_it is not None or pdb_path_it is not None
840
+ ), "Iterable of pdb codes, pdb paths or uniprot IDs required."
841
+
842
+ if pdb_code_it is not None:
843
+ pdbs = pdb_code_it
844
+ source = "pdb_code"
845
+
846
+ if pdb_path_it is not None:
847
+ pdbs = pdb_path_it
848
+ source = "pdb_path"
849
+
850
+ if uniprot_id_it is not None:
851
+ pdbs = uniprot_id_it
852
+ source = "uniprot_id"
853
+
854
+ if chain_selections is None:
855
+ chain_selections = ["all"] * len(pdbs)
856
+
857
+ if model_indices is None:
858
+ model_indices = [1] * len(pdbs)
859
+
860
+ constructor = partial(_mp_graph_constructor, source=source, config=config)
861
+
862
+ graphs = list(
863
+ process_map(
864
+ constructor,
865
+ [
866
+ (pdb, chain_selections[i], model_indices[i])
867
+ for i, pdb in enumerate(pdbs)
868
+ ],
869
+ max_workers=num_cores,
870
+ )
871
+ )
872
+ if out_path is not None:
873
+ [
874
+ nx.write_gpickle(
875
+ g, str(f"{out_path}/" + f"{g.graph['name']}.pickle")
876
+ )
877
+ for g in graphs
878
+ ]
879
+
880
+ if return_dict:
881
+ graphs = {pdb: graphs[i] for i, pdb in enumerate(pdbs)}
882
+
883
+ return graphs
884
+
885
+
886
+ def compute_chain_graph(
887
+ g: nx.Graph,
888
+ chain_list: Optional[List[str]] = None,
889
+ remove_self_loops: bool = False,
890
+ return_weighted_graph: bool = False,
891
+ ) -> Union[nx.Graph, nx.MultiGraph]:
892
+ """Computes a chain-level graph from a protein structure graph.
893
+
894
+ This graph features nodes as individual chains in a complex and edges as
895
+ the interactions between constituent nodes in each chain. You have the
896
+ option of returning an unweighted graph (multigraph,
897
+ ``return_weighted_graph=False``) or a weighted graph
898
+ (``return_weighted_graph=True``). The difference between these is the
899
+ unweighted graph features and edge for each interaction between chains
900
+ (ie the number of edges will be equal to the number of edges in the input
901
+ protein structure graph), while the weighted graph sums these interactions
902
+ to a single edge between chains with the counts stored as features.
903
+
904
+ :param g: A protein structure graph to compute the chain graph of.
905
+ :type g: nx.Graph
906
+ :param chain_list: A list of chains to extract from the input graph.
907
+ If ``None``, all chains will be used. This is provided as input to
908
+ ``extract_subgraph_from_chains``. Default is ``None``.
909
+ :type chain_list: Optional[List[str]]
910
+ :param remove_self_loops: Whether to remove self-loops from the graph.
911
+ Default is False.
912
+ :type remove_self_loops: bool
913
+ :return: A chain-level graph.
914
+ :rtype: Union[nx.Graph, nx.MultiGraph]
915
+ """
916
+ # If we are extracting specific chains, do it here.
917
+ if chain_list is not None:
918
+ g = extract_subgraph_from_chains(g, chain_list)
919
+
920
+ # Initialise new graph with Metadata
921
+ h = nx.MultiGraph()
922
+ h.graph = g.graph
923
+ h.graph["node_type"] = "chain"
924
+
925
+ # Set nodes
926
+ nodes_per_chain = {chain: 0 for chain in g.graph["chain_ids"]}
927
+ sequences = {chain: "" for chain in g.graph["chain_ids"]}
928
+ for n, d in g.nodes(data=True):
929
+ nodes_per_chain[d["chain_id"]] += 1
930
+ sequences[d["chain_id"]] += RESI_THREE_TO_1[d["residue_name"]]
931
+
932
+ h.add_nodes_from(g.graph["chain_ids"])
933
+
934
+ for n, d in h.nodes(data=True):
935
+ d["num_residues"] = nodes_per_chain[n]
936
+ d["sequence"] = sequences[n]
937
+
938
+ # Add edges
939
+ for u, v, d in g.edges(data=True):
940
+ h.add_edge(
941
+ g.nodes[u]["chain_id"], g.nodes[v]["chain_id"], kind=d["kind"]
942
+ )
943
+ # Remove self-loops if necessary. Checks for equality between nodes in a given edge.
944
+ if remove_self_loops:
945
+ edges_to_remove: List[Tuple[str]] = [
946
+ (u, v) for u, v in h.edges() if u == v
947
+ ]
948
+ h.remove_edges_from(edges_to_remove)
949
+
950
+ # Compute a weighted graph if required.
951
+ if return_weighted_graph:
952
+ return compute_weighted_graph_from_multigraph(h)
953
+ return h
954
+
955
+
956
+ def compute_weighted_graph_from_multigraph(g: nx.MultiGraph) -> nx.Graph:
957
+ """Computes a weighted graph from a multigraph.
958
+
959
+ This function is used to convert a multigraph to a weighted graph. The
960
+ weights of the edges are the number of interactions between the nodes.
961
+
962
+ :param g: A multigraph.
963
+ :type g: nx.MultiGraph
964
+ :return: A weighted graph.
965
+ :rtype: nx.Graph
966
+ """
967
+ H = nx.Graph()
968
+ H.graph = g.graph
969
+ H.add_nodes_from(g.nodes(data=True))
970
+ for u, v, d in g.edges(data=True):
971
+ if H.has_edge(u, v):
972
+ H[u][v]["weight"] += len(d["kind"])
973
+ H[u][v]["kind"].update(d["kind"])
974
+ for kind in list(d["kind"]):
975
+ try:
976
+ H[u][v][kind] += 1
977
+ except KeyError:
978
+ H[u][v][kind] = 1
979
+ else:
980
+ H.add_edge(u, v, weight=len(d["kind"]), kind=d["kind"])
981
+ for kind in list(d["kind"]):
982
+ H[u][v][kind] = 1
983
+ return H
984
+
985
+
986
+ def number_groups_of_runs(list_of_values: List[Any]) -> List[str]:
987
+ """Numbers groups of runs in a list of values.
988
+
989
+ E.g. ``["A", "A", "B", "A", "A", "A", "B", "B"] ->
990
+ ["A1", "A1", "B1", "A2", "A2", "A2", "B2", "B2"]``
991
+
992
+ :param list_of_values: List of values to number.
993
+ :type list_of_values: List[Any]
994
+ :return: List of numbered values.
995
+ :rtype: List[str]
996
+ """
997
+ df = pd.DataFrame({"val": list_of_values})
998
+ df["idx"] = df["val"].shift() != df["val"]
999
+ df["sum"] = df.groupby("val")["idx"].cumsum()
1000
+ return list(df["val"].astype(str) + df["sum"].astype(str))
1001
+
1002
+
1003
+ def compute_secondary_structure_graph(
1004
+ g: nx.Graph,
1005
+ allowable_ss_elements: Optional[List[str]] = None,
1006
+ remove_non_ss: bool = True,
1007
+ remove_self_loops: bool = False,
1008
+ return_weighted_graph: bool = False,
1009
+ ) -> Union[nx.Graph, nx.MultiGraph]:
1010
+ """Computes a secondary structure graph from a protein structure graph.
1011
+
1012
+ :param g: A protein structure graph to compute the secondary structure
1013
+ graph of.
1014
+ :type g: nx.Graph
1015
+ :param remove_non_ss: Whether to remove non-secondary structure nodes from
1016
+ the graph. These are denoted as ``"-"`` by DSSP. Default is True.
1017
+ :type remove_non_ss: bool
1018
+ :param remove_self_loops: Whether to remove self-loops from the graph.
1019
+ Default is ``False``.
1020
+ :type remove_self_loops: bool
1021
+ :param return_weighted_graph: Whether to return a weighted graph.
1022
+ Default is False.
1023
+ :type return_weighted_graph: bool
1024
+ :raises ProteinGraphConfigurationError: If the protein structure graph is
1025
+ not configured correctly with secondary structure assignments on all
1026
+ nodes.
1027
+ :return: A secondary structure graph.
1028
+ :rtype: Union[nx.Graph, nx.MultiGraph]
1029
+ """
1030
+ # Initialise list of secondary structure elements we use to build the graph
1031
+ ss_list: List[str] = []
1032
+
1033
+ # Check nodes have secondary structure assignment & store them in list
1034
+ for _, d in g.nodes(data=True):
1035
+ if "ss" not in d.keys():
1036
+ raise ProteinGraphConfigurationError(
1037
+ "Secondary structure not defined for all nodes."
1038
+ )
1039
+ ss_list.append(d["ss"])
1040
+
1041
+ # Number SS elements
1042
+ ss_list = pd.Series(number_groups_of_runs(ss_list))
1043
+ ss_list.index = list(g.nodes())
1044
+
1045
+ # Remove unstructured elements if necessary
1046
+ if remove_non_ss:
1047
+ ss_list = ss_list[~ss_list.str.contains("-")]
1048
+ # Subset to only allowable SS elements if necessary
1049
+ if allowable_ss_elements:
1050
+ ss_list = ss_list[
1051
+ ss_list.str.contains("|".join(allowable_ss_elements))
1052
+ ]
1053
+
1054
+ constituent_residues: Dict[str, List[str]] = ss_list.index.groupby(
1055
+ ss_list.values
1056
+ )
1057
+ constituent_residues = {
1058
+ k: list(v) for k, v in constituent_residues.items()
1059
+ }
1060
+ residue_counts: Dict[str, int] = ss_list.groupby(ss_list).count().to_dict()
1061
+
1062
+ # Add Nodes from secondary structure list
1063
+ h = nx.MultiGraph()
1064
+ h.add_nodes_from(ss_list)
1065
+ nx.set_node_attributes(h, residue_counts, "residue_counts")
1066
+ nx.set_node_attributes(h, constituent_residues, "constituent_residues")
1067
+ # Assign ss
1068
+ for n, d in h.nodes(data=True):
1069
+ d["ss"] = n[0]
1070
+
1071
+ # Add graph-level metadata
1072
+ h.graph = g.graph
1073
+ h.graph["node_type"] = "secondary_structure"
1074
+
1075
+ # Iterate over edges in source graph and add SS-SS edges to new graph.
1076
+ for u, v, d in g.edges(data=True):
1077
+ try:
1078
+ h.add_edge(
1079
+ ss_list[u], ss_list[v], kind=d["kind"], source=f"{u}_{v}"
1080
+ )
1081
+ except KeyError as e:
1082
+ log.debug(
1083
+ f"Edge {u}-{v} not added to secondary structure graph. \
1084
+ Reason: {e} not in graph"
1085
+ )
1086
+
1087
+ # Remove self-loops if necessary.
1088
+ # Checks for equality between nodes in a given edge.
1089
+ if remove_self_loops:
1090
+ edges_to_remove: List[Tuple[str]] = [
1091
+ (u, v) for u, v in h.edges() if u == v
1092
+ ]
1093
+ h.remove_edges_from(edges_to_remove)
1094
+
1095
+ # Create weighted graph from h
1096
+ if return_weighted_graph:
1097
+ return compute_weighted_graph_from_multigraph(h)
1098
+ return h
1099
+
1100
+
1101
+ def compute_line_graph(g: nx.Graph, repopulate_data: bool = True) -> nx.Graph:
1102
+ """Computes the line graph of a graph.
1103
+
1104
+ The line graph of a graph G has a node for each edge in G and an edge
1105
+ joining those nodes if the two edges in G share a common node. For directed
1106
+ graphs, nodes are adjacent exactly when the edges they represent form a
1107
+ directed path of length two.
1108
+
1109
+ The nodes of the line graph are 2-tuples of nodes in the original graph (or
1110
+ 3-tuples for multigraphs, with the key of the edge as the third element).
1111
+
1112
+ :param g: Graph to compute the line graph of.
1113
+ :type g: nx.Graph
1114
+ :param repopulate_data: Whether or not to map node and edge data to edges
1115
+ and nodes of the line graph, defaults to True
1116
+ :type repopulate_data: bool, optional
1117
+ :return: Line graph of g.
1118
+ :rtype: nx.Graph
1119
+ """
1120
+ l_g = nx.generators.line_graph(g)
1121
+ l_g.graph = g.graph
1122
+
1123
+ if repopulate_data:
1124
+ source_edge_data = {(u, v): d for u, v, d in g.edges(data=True)}
1125
+ nx.set_node_attributes(l_g, source_edge_data)
1126
+
1127
+ node_list = {}
1128
+ for u, v, d in l_g.edges(data=True):
1129
+ node_union = u + v
1130
+ for n in node_union:
1131
+ if node_union.count(n) > 1:
1132
+ node_list[(u, v)] = n
1133
+ break
1134
+
1135
+ source_node_data = {k: g.nodes[v] for k, v in node_list.items()}
1136
+ nx.set_edge_attributes(l_g, source_node_data)
1137
+ return l_g
modeling_prot2text.py ADDED
@@ -0,0 +1,392 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import GPT2Config, AutoTokenizer, GPT2Config
2
+ from transformers import PretrainedConfig, PreTrainedModel
3
+ import transformers
4
+ from typing import Optional, Tuple, Callable
5
+ import torch
6
+ import torch.nn as nn
7
+ from transformers.modeling_utils import PreTrainedModel, PretrainedConfig
8
+ from .utils import CABlock, _GPT2LMHeadModel
9
+ from .configuration_prot2text import Prot2TextConfig
10
+ import os
11
+ import numpy as np
12
+ from transformers.generation.configuration_utils import GenerationConfig
13
+ from transformers.generation.logits_process import LogitsProcessorList
14
+ from transformers.generation.stopping_criteria import StoppingCriteriaList
15
+
16
+ from .pdb2graph import PDB2Graph, download_alphafold_structure
17
+ from .graphs import *
18
+ from .utils_dataset import *
19
+
20
+ from graphein.protein.config import ProteinGraphConfig, DSSPConfig
21
+ from graphein.protein.features.nodes.amino_acid import amino_acid_one_hot, meiler_embedding, expasy_protein_scale, hydrogen_bond_acceptor, hydrogen_bond_donor
22
+ from graphein.protein.features.nodes.dssp import phi, psi, asa, rsa, secondary_structure
23
+ from graphein.protein.edges.distance import (add_peptide_bonds,
24
+ add_hydrogen_bond_interactions,
25
+ add_distance_threshold,
26
+ )
27
+
28
+ from torch_geometric.nn import RGCNConv, global_mean_pool
29
+
30
+
31
+
32
+ class EncoderRGCN(PreTrainedModel):
33
+ '''
34
+ This class implement the RGCN encoder to encode the protein structure
35
+ '''
36
+ def __init__(self, input_dim, hidden_dim=512, n_layers=6, emb_dim=512, dropout=0.2, num_relation=7, prot2text_version='1.0'):
37
+ super(EncoderRGCN, self).__init__(PretrainedConfig(name='RGCN'))
38
+ self.n_layers = n_layers
39
+ self.output_dim = emb_dim
40
+ self.prot2text_version = prot2text_version
41
+
42
+ self.fc0 = nn.Linear(input_dim, hidden_dim)
43
+ self.batchnorm_final = nn.BatchNorm1d(hidden_dim)
44
+
45
+ self.batch_norms = nn.ModuleList()
46
+ self.batch_norms.append(nn.BatchNorm1d(hidden_dim))
47
+ lst = list()
48
+
49
+ lst.append(RGCNConv(hidden_dim, hidden_dim, num_relations=num_relation))
50
+
51
+ for i in range(n_layers-1):
52
+ lst.append(RGCNConv(hidden_dim,hidden_dim, num_relations=num_relation))
53
+
54
+ self.conv = nn.ModuleList(lst)
55
+
56
+ self.fc1 = nn.Linear(hidden_dim, hidden_dim)
57
+ self.fc2 = nn.Linear(hidden_dim, self.output_dim)
58
+
59
+ self.dropout = nn.Dropout(p=dropout)
60
+ self.relu = nn.LeakyReLU()
61
+ self.batchnorm = nn.BatchNorm1d(hidden_dim)
62
+ self.main_input_name = 'nothing'
63
+
64
+ def forward(self, x:Optional[torch.FloatTensor] = None,
65
+ edge_index:Optional[torch.LongTensor] = None,
66
+ edge_type:Optional[torch.LongTensor] = None,
67
+ batch:Optional[torch.LongTensor] = None,
68
+ **kargs):
69
+ #construct pyg edge index shape (2, num_edges) from edge_list
70
+ x = self.relu(self.fc0(x))
71
+
72
+ for i in range(self.n_layers):
73
+ x = self.conv[i](x, edge_index, edge_type)
74
+
75
+ out = global_mean_pool(x, batch)
76
+ out = self.relu(self.fc1(out))
77
+ out = self.relu(self.fc2(out))
78
+
79
+ return out.unsqueeze(1)
80
+
81
+ class Prot2TextModel(PreTrainedModel):
82
+ config_class = Prot2TextConfig
83
+ _keys_to_ignore_on_load_missing = [r"transformer"]
84
+ base_model_prefix = "decoder"
85
+ def __init__(self, config):
86
+ super().__init__(config)
87
+
88
+ self.gpt_config = GPT2Config.from_dict(config.gpt_config)
89
+
90
+ # if we are using RGCN to encode the protein's structure, define the RGCN encoder
91
+ if config.rgcn:
92
+ self.encoder = EncoderRGCN(input_dim=config.rgcn_input_dim, hidden_dim=self.gpt_config.n_embd, n_layers=config.rgcn_n_layers, emb_dim=self.gpt_config.n_embd, prot2text_version=self.config.prot2text_version)
93
+
94
+ # define the GPT2 decoder
95
+ self.decoder = _GPT2LMHeadModel(self.gpt_config)
96
+
97
+ # if using ESM to encode protein's sequence, define the ESM layer, the Projection layer and the fusion layer
98
+ if config.esm:
99
+ self.esm_config = PretrainedConfig.from_dict(config.esm_config)
100
+ self.esm = transformers.EsmModel(self.esm_config)
101
+ self.to_embedding = nn.Linear(self.esm_config.hidden_size, self.gpt_config.n_embd)
102
+ if config.cross_esm_graph and config.rgcn:
103
+ self.h = nn.ModuleList([CABlock(self.gpt_config, layer_idx=i) for i in range(4)])
104
+ self.ln_f = nn.LayerNorm(self.gpt_config.n_embd, eps=self.gpt_config.layer_norm_epsilon)
105
+
106
+ self.config = config
107
+
108
+
109
+ def get_encoder(self):
110
+ return self.encoder
111
+
112
+ def get_decoder(self):
113
+ return self.decoder
114
+
115
+ def get_input_embeddings(self):
116
+ if hasattr(self, "transformer"):
117
+ return self.transformer.wte
118
+ return self.decoder.transformer.wte
119
+
120
+ def warm_up(self, gpt_model=None, esm_model=None):
121
+ if esm_model is not None:
122
+ self.esm = transformers.EsmModel.from_pretrained(esm_model)
123
+ if gpt_model is not None:
124
+ self.decoder = _GPT2LMHeadModel.from_pretrained(gpt_model, add_cross_attention=True, use_cache=False)
125
+ self.decoder.resize_token_embeddings(self.gpt_config.vocab_size)
126
+ self.decoder.config = self.gpt_config
127
+
128
+
129
+ def forward(self,
130
+ encoder_input_ids: Optional[torch.LongTensor] = None,
131
+ edge_index: Optional[torch.LongTensor] = None,
132
+ batch: Optional[torch.LongTensor] = None,
133
+ x: Optional[torch.FloatTensor] = None,
134
+ edge_type: Optional[torch.LongTensor] = None,
135
+ decoder_input_ids: Optional[torch.LongTensor] = None,
136
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
137
+ past_key_values_graph_esm: Optional[Tuple[Tuple[torch.Tensor]]] = None,
138
+ decoder_attention_mask: Optional[torch.FloatTensor] = None,
139
+ attention_mask: Optional[torch.FloatTensor] = None,
140
+ token_type_ids: Optional[torch.LongTensor] = None,
141
+ position_ids: Optional[torch.LongTensor] = None,
142
+ head_mask: Optional[torch.FloatTensor] = None,
143
+ inputs_embeds: Optional[torch.FloatTensor] = None,
144
+ encoder_hidden_states: Optional[torch.Tensor] = None,
145
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
146
+ labels: Optional[torch.LongTensor] = None,
147
+ use_cache: Optional[bool] = None,
148
+ output_attentions: Optional[bool] = None,
149
+ output_hidden_states: Optional[bool] = None,
150
+ return_dict: Optional[bool] = None,
151
+ get_graph_emb: Optional[bool] = False,
152
+ **delete_args,
153
+ ):
154
+ use_cache = use_cache if use_cache is not None else self.gpt_config.use_cache
155
+ return_dict = return_dict if return_dict is not None else self.gpt_config.use_return_dict
156
+
157
+
158
+ if decoder_input_ids is not None and len(decoder_input_ids.size()) == 3:
159
+ decoder_input_ids = decoder_input_ids.squeeze(0)
160
+
161
+ if x is not None and self.config.rgcn:
162
+ graph_emb = self.encoder(x, edge_index, edge_type, batch)
163
+ graph_mask = None
164
+
165
+ if self.config.esm:
166
+ if self.config.prot2text_version=='1.0':
167
+ if encoder_input_ids.size()[1] != 1021:
168
+ raise ValueError("For this version of the model you need to PAD/Truncate the amino acid sequence for the ESM model to 1021")
169
+
170
+ esm_emb = self.esm(input_ids=encoder_input_ids, attention_mask=attention_mask, return_dict=return_dict).last_hidden_state
171
+ esm_emb = self.to_embedding(esm_emb)
172
+ if not self.config.cross_esm_graph and self.config.rgcn:
173
+ graph_emb = torch.cat((graph_emb, esm_emb), dim=1)
174
+ t_add = torch.ones((attention_mask.size(0), 1)).to(attention_mask.get_device())
175
+ attention_mask = torch.cat((t_add, attention_mask), dim=1)
176
+ elif self.config.cross_esm_graph and self.config.rgcn:
177
+ if past_key_values_graph_esm is None:
178
+ past_length = 0
179
+ past_key_values_graph_esm = tuple([None] * len(self.h))
180
+ else:
181
+ past_length = past_key_values_graph_esm[0][0].size(-2)
182
+ output_shape = esm_emb.size()
183
+
184
+ all_self_attentions = () if output_attentions else None
185
+ all_cross_attentions = () if output_attentions and self.gpt_config.add_cross_attention else None
186
+ all_hidden_states = () if output_hidden_states else None
187
+ for i, (block, layer_past) in enumerate(zip(self.h, past_key_values_graph_esm)):
188
+ outputs = block(
189
+ esm_emb,
190
+ layer_past=layer_past,
191
+ attention_mask=attention_mask,
192
+ encoder_hidden_states=graph_emb,
193
+ encoder_attention_mask=graph_mask,
194
+ use_cache=use_cache,
195
+ output_attentions=False,
196
+ )
197
+ esm_emb = outputs[0]
198
+
199
+ esm_emb = self.ln_f(esm_emb)
200
+ esm_emb = esm_emb.view(output_shape)
201
+ graph_emb = esm_emb
202
+ else:
203
+ graph_emb = esm_emb
204
+ else:
205
+ attention_mask = None
206
+ if self.config.prot2text_version=='1.0':
207
+ attention_mask = None
208
+ if get_graph_emb:
209
+ return graph_emb
210
+
211
+ transformer_outputs = self.decoder(input_ids=decoder_input_ids,
212
+ past_key_values=past_key_values,
213
+ attention_mask=decoder_attention_mask,
214
+ token_type_ids=token_type_ids,
215
+ position_ids=position_ids,
216
+ head_mask=head_mask,
217
+ inputs_embeds=inputs_embeds,
218
+ encoder_hidden_states=graph_emb,
219
+ encoder_attention_mask=attention_mask,
220
+ labels=labels,
221
+ use_cache=use_cache,
222
+ output_attentions=output_attentions,
223
+ output_hidden_states=output_hidden_states,
224
+ return_dict=return_dict,
225
+ )
226
+
227
+ return transformer_outputs
228
+
229
+ @torch.no_grad()
230
+ def generate_protein_description(self,
231
+ protein_pdbID=None,
232
+ protein_sequence=None,
233
+ edge_index: Optional[torch.LongTensor] = None,
234
+ x: Optional[torch.FloatTensor] = None,
235
+ edge_type: Optional[torch.LongTensor] = None,
236
+ tokenizer=None,
237
+ device='cpu'
238
+ ):
239
+
240
+ if self.config.esm and not self.config.rgcn and protein_sequence==None:
241
+ raise ValueError(
242
+ "The model you are trying to use is based only on protein sequence, please provide an amino-acid protein_sequence"
243
+ )
244
+ if self.config.rgcn and protein_pdbID==None and (x==None or edge_index==None or edge_type==None):
245
+ raise ValueError(
246
+ "The model you are trying to use is based on protein structure, please provide a AlphaFold ID (you must have to have internet connection using protein_pdbID, or provide the triplet inputs: x (node features), edge_index and edge_type"
247
+ )
248
+ if self.config.esm:
249
+ esmtokenizer = AutoTokenizer.from_pretrained(self.config.esm_model_name)
250
+
251
+ if protein_pdbID==None and protein_sequence==None:
252
+ raise ValueError(
253
+ "you need to provide either a protein AlphaFold Id or an amino-acid sequence"
254
+ )
255
+
256
+ if protein_pdbID!=None:
257
+ config = {"node_metadata_functions": [amino_acid_one_hot,
258
+ expasy_protein_scale,
259
+ meiler_embedding,
260
+ hydrogen_bond_acceptor, hydrogen_bond_donor
261
+ ],
262
+ "edge_construction_functions": [add_peptide_bonds,
263
+ add_hydrogen_bond_interactions,
264
+ partial(add_distance_threshold, long_interaction_threshold=3, threshold=10.),],
265
+ "graph_metadata_functions":[asa,phi, psi, secondary_structure, rsa],
266
+ "dssp_config": DSSPConfig()}
267
+ config = ProteinGraphConfig(**config)
268
+
269
+ PATH_TO_DATA = f"~/.tmp/pdb/pdb"
270
+ OUTPUT_FOLDER = f"~/.tmp/pdb/raw"
271
+ save_dir = f"~/.tmp/pdb/"
272
+ isExist = os.path.exists(PATH_TO_DATA)
273
+ if not isExist:
274
+ os.makedirs(PATH_TO_DATA)
275
+ isExist = os.path.exists(OUTPUT_FOLDER)
276
+ if not isExist:
277
+ os.makedirs(OUTPUT_FOLDER)
278
+ isExist = os.path.exists(save_dir+'processed')
279
+ if not isExist:
280
+ os.makedirs(save_dir+'processed')
281
+
282
+ structure_filename = download_alphafold_structure(uniprot_id=protein_pdbID, out_dir=PATH_TO_DATA)
283
+ if structure_filename is None:
284
+ raise ValueError("Error! the ID does not exist in AlphaFoldDB or you do not have internet connection")
285
+ graph_filename = structure_filename.split('/')
286
+ graph_filename[-2] = 'raw'
287
+ graph_filename[-1] = graph_filename[-1].replace('.pdb', '.pt')
288
+ graph_filename = '/'.join(graph_filename)
289
+ process_filename = structure_filename.split('/')
290
+ process_filename[-2] = 'processed'
291
+ process_filename[-1] = process_filename[-1].replace('.pdb', '.pt')
292
+ process_filename = '/'.join(process_filename)
293
+ try:
294
+ gpdb = PDB2Graph(root = PATH_TO_DATA, output_folder = OUTPUT_FOLDER, config=config, n_processors=1).create_pyg_graph(structure_filename)
295
+ seq = esmtokenizer(gpdb.sequence, add_special_tokens=True, truncation=True, max_length=1021, padding='max_length',return_tensors="pt") #
296
+ torch.save(gpdb, graph_filename)
297
+ gpdb.edge_type = [np.array(gpdb.edge_type.transpose(0,1))]
298
+ gpdb.encoder_input_ids = seq['input_ids']
299
+ gpdb.attention_mask = seq['attention_mask']
300
+ torch.save(gpdb, process_filename)
301
+ except:
302
+ os.remove(structure_filename)
303
+ raise ValueError('creating graphs did not work, probably the pdb file of alphaFold is damaged')
304
+
305
+ self.eval()
306
+ inputs = gpdb
307
+ inputs = inputs.to_dict()
308
+
309
+ inputs['edge_type'] = torch.cat([torch.tensor(inputs['edge_type'][i]) for i in range(len(inputs['edge_type']))], dim=0)
310
+ inputs['edge_type'] = torch.argmax(inputs['edge_type'], dim=1)
311
+ for key in ['num_nodes', 'node_id', 'name', 'sequence', 'distance_matrix', 'distance', 'coordinates']:
312
+ inputs.pop(key)
313
+ inputs['decoder_input_ids'] = inputs['encoder_input_ids'][:,0:1].clone()
314
+ inputs['decoder_input_ids'][:,0] = tokenizer.bos_token_id
315
+ inputs["decoder_attention_mask"] = torch.ones(inputs['decoder_input_ids'].shape[0], 1)
316
+ self.to(device)
317
+ inputs = {k: v.to(device=device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
318
+ encoder_state = dict()
319
+ encoder_state['hidden_states'] = self(**inputs, get_graph_emb=True, output_attentions=True)
320
+ encoder_state['attentions'] = inputs['attention_mask']
321
+ for key in ['edge_index', 'edge_type', 'x', 'encoder_input_ids']:
322
+ inputs.pop(key)
323
+ tok_ids = self.decoder.generate(input_ids=inputs['decoder_input_ids'],
324
+ encoder_outputs=encoder_state,
325
+ use_cache=True,
326
+ output_attentions=False,
327
+ output_scores=False,
328
+ return_dict_in_generate=True,
329
+ encoder_attention_mask=inputs['attention_mask'],
330
+ length_penalty=1.0,
331
+ no_repeat_ngram_size=None,
332
+ early_stopping=False,
333
+ num_beams=1)
334
+
335
+ generated = tokenizer.batch_decode(tok_ids.get('sequences'), skip_special_tokens=True)
336
+
337
+ os.remove(structure_filename)
338
+ os.remove(graph_filename)
339
+ os.remove(process_filename)
340
+
341
+ return generated[0].replace('<|stop_token|>', '').replace('<|graph_token|>', '')
342
+
343
+ else:
344
+ seq = esmtokenizer([protein_sequence], add_special_tokens=True, truncation=True, max_length=1021, padding='max_length', return_tensors="pt")
345
+ inputs={}
346
+ inputs['encoder_input_ids'] = seq['input_ids']
347
+ inputs['attention_mask'] = seq['attention_mask']
348
+ inputs['decoder_input_ids'] = inputs['encoder_input_ids'][:,0:1].clone()
349
+ inputs['decoder_input_ids'][:,0] = tokenizer.bos_token_id
350
+
351
+ self.to(device)
352
+ inputs = {k: v.to(device=device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
353
+ encoder_state = dict()
354
+ encoder_state['hidden_states'] = self(**inputs, get_graph_emb=True, output_attentions=True)
355
+ generated = tokenizer.batch_decode(self.decoder.generate(input_ids=inputs['decoder_input_ids'], encoder_outputs=encoder_state, use_cache=True), skip_special_tokens=True)
356
+
357
+ return generated[0].replace('<|stop_token|>', '').replace('<|graph_token|>', '')
358
+
359
+ @torch.no_grad()
360
+ def generate(self,
361
+ inputs: Optional[torch.Tensor] = None,
362
+ generation_config: Optional[GenerationConfig] = None,
363
+ logits_processor: Optional[LogitsProcessorList] = None,
364
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
365
+ prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
366
+ synced_gpus: Optional[bool] = None,
367
+ assistant_model: Optional["PreTrainedModel"] = None,
368
+ streamer: Optional["BaseStreamer"] = None,
369
+ **kwargs,
370
+ ):
371
+ encoder_state = self(**kwargs, get_graph_emb=True)
372
+ input_ids = kwargs['decoder_input_ids']
373
+ attention_mask = kwargs['decoder_attention_mask']
374
+ kwargs['encoder_attention_mask'] = kwargs['attention_mask']
375
+ if not self.config.cross_esm_graph and self.config.rgcn and self.config.esm:
376
+ t_add = torch.ones((kwargs['encoder_attention_mask'].size(0), 1)).to(kwargs['encoder_attention_mask'].get_device())
377
+ kwargs['encoder_attention_mask'] = torch.cat((t_add, kwargs['encoder_attention_mask']), dim=1)
378
+ for key in ['edge_index', 'edge_type', 'x', 'encoder_input_ids', 'decoder_input_ids', 'decoder_attention_mask', 'batch', 'attention_mask', 'max_length',
379
+ '_num_nodes', 'node_id', 'name', 'sequence', 'distance_matrix', 'distance', 'coordinates', 'ptr', 'num_nodes',]:
380
+ if key in kwargs.keys():
381
+ kwargs.pop(key)
382
+ return self.decoder.generate(input_ids=input_ids,
383
+ generation_config=generation_config,
384
+ logits_processor=logits_processor,
385
+ stopping_criteria=stopping_criteria,
386
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
387
+ synced_gpus=synced_gpus,
388
+ assistant_model=assistant_model,
389
+ streamer=streamer,
390
+ encoder_outputs={'hidden_states': encoder_state, 'attentions':0},
391
+ **kwargs
392
+ )
pdb2graph.py ADDED
@@ -0,0 +1,171 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import multiprocessing
2
+ import os
3
+ from tqdm import tqdm
4
+ from sklearn.preprocessing import MultiLabelBinarizer
5
+
6
+ from torch_geometric.data import Data
7
+ import torch
8
+
9
+ import numpy as np
10
+
11
+ from .conversion import convert_nx_to_pyg_data
12
+ from graphein.protein.config import ProteinGraphConfig, DSSPConfig
13
+ from graphein.protein.features.nodes.amino_acid import amino_acid_one_hot, meiler_embedding, expasy_protein_scale, hydrogen_bond_acceptor, hydrogen_bond_donor
14
+ from graphein.protein.features.nodes.dssp import phi, psi, asa, rsa, secondary_structure
15
+ from graphein.protein.edges.distance import (add_peptide_bonds,
16
+ add_hydrogen_bond_interactions,
17
+ add_disulfide_interactions,
18
+ add_ionic_interactions,
19
+ add_delaunay_triangulation,
20
+ add_distance_threshold,
21
+ add_sequence_distance_edges,
22
+ add_k_nn_edges)
23
+
24
+ from functools import partial
25
+ from .graphs import *
26
+ from .utils_dataset import *
27
+ import os
28
+ import sys
29
+ import subprocess
30
+ import wget
31
+
32
+
33
+ class PDB2Graph():
34
+ def __init__(self, root, output_folder, config, n_processors=int(multiprocessing.cpu_count())):
35
+ self.root = root
36
+ self.output_folder = output_folder
37
+ self.map_secondary_structure = {'-':0, 'H':1, 'B':2, 'E':3, 'G':4, 'I':5, 'T':6, 'S':7}
38
+ self.init_ohe_edge_type()
39
+ self.config = config
40
+ self.features = ['phi', 'psi', 'rsa', 'asa', 'ss', 'expasy']
41
+ self.n_processors = n_processors
42
+ self.raw_dir = root
43
+ self.processed_dir = self._processed_dir()
44
+ self.raw_file_names = self._raw_file_names()
45
+ self.processed_file_names = self._processed_file_names()
46
+
47
+
48
+ def _processed_dir(self):
49
+ #processed_dir = os.path.join(os.path.split(self.root)[0], "processed_new")
50
+ if not os.path.exists(self.output_folder):
51
+ os.makedirs(self.output_folder)
52
+ return self.output_folder
53
+
54
+ def _raw_file_names(self):
55
+ return os.listdir(self.raw_dir)
56
+
57
+ def _processed_file_names(self):
58
+ return [self.pdb2pathdata(pdb_path.split(".")[0]) for pdb_path in self.raw_file_names]
59
+
60
+ def create_nx_graph(self, path_to_structure):
61
+ return construct_graph(self.config, pdb_path = path_to_structure)
62
+
63
+ def create_pyg_graph(self, path_to_structure):
64
+ pyg_graph = convert_nx_to_pyg_data(self.create_nx_graph(path_to_structure))
65
+
66
+ graph = Data(edge_index = pyg_graph.edge_index,
67
+ num_nodes = len(pyg_graph.node_id),
68
+ node_id = pyg_graph.node_id,
69
+ name = pyg_graph.name[0],
70
+ sequence = getattr(pyg_graph, f"sequence_{pyg_graph.chain_id[0]}"),
71
+ distance_matrix = pyg_graph.dist_mat,
72
+ distance = pyg_graph.distance,
73
+ coordinates = torch.FloatTensor(np.array(pyg_graph.coords[0])))
74
+ #create the features
75
+ x = np.array([np.argmax(pyg_graph.amino_acid_one_hot, axis=1)]).reshape(-1,1)
76
+ for feat in self.features:
77
+ if feat == "ss":
78
+ feature = np.array([[self.map_secondary_structure.get(feat_node, 0)] \
79
+ for feat_node in pyg_graph[feat]])
80
+ else:
81
+ feature = np.array(pyg_graph[feat])
82
+ if len(feature.shape) == 1:
83
+ feature = feature.reshape(-1,1)
84
+ x = np.concatenate((x, feature), axis = 1)
85
+ graph.edge_type = self.mlb.transform(pyg_graph.kind)
86
+ graph.x = torch.FloatTensor(x)
87
+ # y = self.annotations[graph.name.split("_")[0]]
88
+ # if self.task == 'GeneOntology' :
89
+ # graph.y_mf = torch.FloatTensor(y["mf"])
90
+ # graph.y_cc = torch.FloatTensor(y["cc"])
91
+ # graph.y_bp = torch.FloatTensor(y["bp"])
92
+ # else:
93
+ # graph.y_ec = torch.FloatTensor(y["ec"])
94
+ return graph
95
+
96
+ def init_ohe_edge_type(self):
97
+ self.mlb = MultiLabelBinarizer(classes = ['peptide_bond', 'sequence_distance_2', 'sequence_distance_3'
98
+ , 'distance_threshold', 'delaunay', 'hbond', 'k_nn'])
99
+ self.mlb.fit([['peptide_bond', 'sequence_distance_2', 'sequence_distance_3'
100
+ , 'distance_threshold', 'delaunay', 'hbond', 'k_nn']])
101
+
102
+ def process(self):
103
+ """Convert the PDB files into torch geometric graphs"""
104
+ # self.pdb2graph = PDB2Graph(self.config)
105
+ to_be_processed = self.get_files_to_process()
106
+
107
+ # pool = multiprocessing.Pool(self.n_processors)
108
+ # for _ in tqdm(pool.imap_unordered(self.graph_creation, to_be_processed), total=len(to_be_processed)):
109
+ # continue
110
+ # pool.close()
111
+ # pool.join()
112
+
113
+
114
+
115
+ processes = []
116
+ for prot in tqdm(to_be_processed):
117
+ p = multiprocessing.Process(target=self.graph_creation, args=(prot,))
118
+ processes.append(p)
119
+ p.start()
120
+
121
+ for process in processes:
122
+ process.join()
123
+
124
+
125
+ def graph_creation(self, pdb):
126
+ """Create a graph from the PDB file"""
127
+
128
+ # Define the path_to_structure from the pdb name file
129
+ path_to_structure = self.pdb2pathstructure(pdb)
130
+
131
+ # Convert the structure into a graph
132
+ g = self.create_pyg_graph(path_to_structure)
133
+ # Save the graph
134
+ torch.save(g, os.path.join(self.output_folder, self.pdb2pathdata(pdb)))
135
+
136
+ return None
137
+
138
+ def pdb2pathdata(self, pdb):
139
+ return pdb+'.pt'
140
+
141
+ def pdb2pathstructure(self, pdb):
142
+ return os.path.join(self.raw_dir, pdb+'.pdb')
143
+
144
+ def get_files_to_process(self):
145
+ RAW_FILES = self.processed_file_names
146
+ PROCESSED_FILES = os.listdir(self.processed_dir)
147
+ to_be_processed = set(RAW_FILES).difference(set(PROCESSED_FILES))
148
+ to_be_processed = [path.split('.')[0] for path in to_be_processed]
149
+ return to_be_processed
150
+
151
+ def download_alphafold_structure(
152
+ uniprot_id: str,
153
+ out_dir: str,
154
+ version: int = 4
155
+ ):
156
+
157
+ BASE_URL = "https://alphafold.ebi.ac.uk/files/"
158
+ uniprot_id = uniprot_id.upper()
159
+
160
+ query_url = f"{BASE_URL}AF-{uniprot_id}-F1-model_v{version}.pdb"
161
+ structure_filename = os.path.join(out_dir, f"AF-{uniprot_id}-F1-model_v{version}.pdb")
162
+ if os.path.exists(structure_filename):
163
+ return structure_filename
164
+ try:
165
+ structure_filename = wget.download(query_url, out=out_dir)
166
+ except:
167
+ print('Error.. could not download: ', f"AF-{uniprot_id}-F1-model_v{version}.pdb")
168
+ return None
169
+ return structure_filename
170
+
171
+
utils.py ADDED
@@ -0,0 +1,742 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch.nn as nn
2
+ from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2MLP
3
+ from typing import Optional, Tuple, Union, Any, Dict, List
4
+ from transformers import Seq2SeqTrainer, GPT2LMHeadModel
5
+ from torch.utils.data.distributed import DistributedSampler
6
+ import torch
7
+ from transformers.deepspeed import is_deepspeed_zero3_enabled
8
+ from transformers.generation.logits_process import LogitsProcessorList
9
+ from transformers.generation.stopping_criteria import StoppingCriteriaList
10
+ from transformers.generation.utils import GreedySearchOutput, GreedySearchEncoderDecoderOutput, BeamSearchOutput, BeamSearchEncoderDecoderOutput
11
+ from transformers.generation.beam_search import BeamScorer
12
+
13
+ from torch_geometric.loader import DataLoader
14
+ from torch_geometric.data import Dataset
15
+
16
+ class _GPT2LMHeadModel(GPT2LMHeadModel):
17
+ def _init_(self, config):
18
+ super(GPT2LMHeadModel, self).init_(config)
19
+ self.config = config
20
+
21
+
22
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, encoder_outputs=None, **kwargs):
23
+ '''
24
+ This function is an edited version of the prepare_inputs_for_generation function from HuggingFace's transformers
25
+ https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
26
+ '''
27
+ token_type_ids = kwargs.get("token_type_ids", None)
28
+ # only last token for inputs_ids if past is defined in kwargs
29
+ if past_key_values:
30
+ input_ids = input_ids[:, -1].unsqueeze(-1)
31
+ if token_type_ids is not None:
32
+ token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
33
+
34
+ attention_mask = kwargs.get("attention_mask", None)
35
+ position_ids = kwargs.get("position_ids", None)
36
+ if self.config.prot2text_version=="1.1" or self.config.prot2text_version=="1.2":
37
+ encoder_attention_mask = kwargs.get("encoder_attention_mask", None)
38
+ elif self.config.prot2text_version=="1.0":
39
+ encoder_attention_mask = None
40
+
41
+ if attention_mask is not None and position_ids is None:
42
+ position_ids = attention_mask.long().cumsum(-1) - 1
43
+ position_ids.masked_fill_(attention_mask == 0, 1)
44
+ if past_key_values:
45
+ position_ids = position_ids[:, -1].unsqueeze(-1)
46
+ else:
47
+ position_ids = None
48
+
49
+ model_specific_kwargs = {
50
+ "encoder_hidden_states": encoder_outputs['hidden_states'],
51
+ }
52
+
53
+ return {
54
+ "input_ids": input_ids,
55
+ "past_key_values": past_key_values,
56
+ "use_cache": kwargs.get("use_cache"),
57
+ "position_ids": position_ids,
58
+ "attention_mask": attention_mask,
59
+ "token_type_ids": token_type_ids,
60
+ "encoder_attention_mask": encoder_attention_mask,
61
+ **model_specific_kwargs
62
+ }
63
+
64
+
65
+ def greedy_search(
66
+ self,
67
+ input_ids: torch.LongTensor,
68
+ logits_processor: Optional[LogitsProcessorList] = None,
69
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
70
+ max_length: Optional[int] = None,
71
+ pad_token_id: Optional[int] = None,
72
+ eos_token_id: Optional[Union[int, List[int]]] = None,
73
+ output_attentions: Optional[bool] = None,
74
+ output_hidden_states: Optional[bool] = None,
75
+ output_scores: Optional[bool] = None,
76
+ return_dict_in_generate: Optional[bool] = None,
77
+ synced_gpus: bool = False,
78
+ streamer: Optional["BaseStreamer"] = None,
79
+ **model_kwargs,
80
+ ) -> Union[GreedySearchOutput, torch.LongTensor]:
81
+ '''
82
+ This function is an edited version of the greedy_search function from HuggingFace's transformers
83
+ https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py
84
+ '''
85
+
86
+ # init values
87
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
88
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
89
+ if max_length is not None:
90
+ warnings.warn(
91
+ "`max_length` is deprecated in this function, use"
92
+ " `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.",
93
+ UserWarning,
94
+ )
95
+ stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
96
+ pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
97
+ eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
98
+ if isinstance(eos_token_id, int):
99
+ eos_token_id = [eos_token_id]
100
+ eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
101
+ output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
102
+ output_attentions = (
103
+ output_attentions if output_attentions is not None else self.generation_config.output_attentions
104
+ )
105
+ output_hidden_states = (
106
+ output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
107
+ )
108
+ return_dict_in_generate = (
109
+ return_dict_in_generate
110
+ if return_dict_in_generate is not None
111
+ else self.generation_config.return_dict_in_generate
112
+ )
113
+
114
+ # init attention / hidden states / scores tuples
115
+ scores = () if (return_dict_in_generate and output_scores) else None
116
+ decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
117
+ cross_attentions = () if (return_dict_in_generate and output_attentions) else None
118
+ decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
119
+
120
+ # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
121
+ if return_dict_in_generate and self.config.is_encoder_decoder:
122
+ encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
123
+ encoder_hidden_states = (
124
+ model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
125
+ )
126
+
127
+ # keep track of which sequences are already finished
128
+ unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
129
+
130
+ this_peer_finished = False # used by synced_gpus only
131
+ while True:
132
+ if synced_gpus:
133
+ # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
134
+ # The following logic allows an early break if all peers finished generating their sequence
135
+ this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
136
+ # send 0.0 if we finished, 1.0 otherwise
137
+ dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
138
+ # did all peers finish? the reduced sum will be 0.0 then
139
+ if this_peer_finished_flag.item() == 0.0:
140
+ break
141
+
142
+ # prepare model inputs
143
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
144
+
145
+ # forward pass to get next token
146
+ outputs = self(
147
+ **model_inputs,
148
+ return_dict=True,
149
+ output_attentions=output_attentions,
150
+ output_hidden_states=output_hidden_states,
151
+ )
152
+
153
+ if synced_gpus and this_peer_finished:
154
+ continue # don't waste resources running the code we don't need
155
+
156
+ next_token_logits = outputs.logits[:, -1, :]
157
+
158
+ # pre-process distribution
159
+ next_tokens_scores = logits_processor(input_ids, next_token_logits)
160
+
161
+ # Store scores, attentions and hidden_states when required
162
+ if return_dict_in_generate:
163
+ if output_scores:
164
+ scores += (next_tokens_scores,)
165
+ if output_attentions:
166
+ decoder_attentions += (
167
+ (outputs.decoder_attentions,) if not self.config.is_encoder_decoder else (outputs.attentions,)
168
+ )
169
+ if self.config.is_encoder_decoder:
170
+ cross_attentions += (outputs.cross_attentions,)
171
+
172
+ if output_hidden_states:
173
+ decoder_hidden_states += (
174
+ (outputs.decoder_hidden_states,)
175
+ if self.config.is_encoder_decoder
176
+ else (outputs.hidden_states,)
177
+ )
178
+
179
+ # argmax
180
+ next_tokens = torch.argmax(next_tokens_scores, dim=-1)
181
+
182
+ # finished sentences should have their next token be a padding token
183
+ if eos_token_id is not None:
184
+ if pad_token_id is None:
185
+ raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
186
+ next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
187
+
188
+ # update generated ids, model inputs, and length for next step
189
+ input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
190
+ if streamer is not None:
191
+ streamer.put(next_tokens.cpu())
192
+ model_kwargs = self._update_model_kwargs_for_generation(
193
+ outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
194
+ )
195
+
196
+ # if eos_token was found in one sentence, set sentence to finished
197
+ if eos_token_id_tensor is not None:
198
+ unfinished_sequences = unfinished_sequences.mul(
199
+ next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
200
+ )
201
+
202
+ # stop when each sentence is finished
203
+ if unfinished_sequences.max() == 0:
204
+ this_peer_finished = True
205
+
206
+ # stop if we exceed the maximum length
207
+ try:
208
+ if stopping_criteria(input_ids, scores):
209
+ this_peer_finished = True
210
+ except:
211
+ if all(stopping_criteria(input_ids, scores)):
212
+ this_peer_finished = True
213
+
214
+ if this_peer_finished and not synced_gpus:
215
+ break
216
+
217
+ if streamer is not None:
218
+ streamer.end()
219
+
220
+ if return_dict_in_generate:
221
+ if self.config.is_encoder_decoder:
222
+ return GreedySearchEncoderDecoderOutput(
223
+ sequences=input_ids,
224
+ scores=scores,
225
+ encoder_attentions=encoder_attentions,
226
+ encoder_hidden_states=encoder_hidden_states,
227
+ decoder_attentions=decoder_attentions,
228
+ cross_attentions=cross_attentions,
229
+ decoder_hidden_states=decoder_hidden_states,
230
+ )
231
+ else:
232
+ return GreedySearchDecoderOnlyOutput(
233
+ sequences=input_ids,
234
+ scores=scores,
235
+ attentions=decoder_attentions,
236
+ hidden_states=decoder_hidden_states,
237
+ )
238
+ else:
239
+ return input_ids
240
+
241
+ def _greedy_search(
242
+ self,
243
+ input_ids: torch.LongTensor,
244
+ logits_processor: Optional[LogitsProcessorList] = None,
245
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
246
+ max_length: Optional[int] = None,
247
+ pad_token_id: Optional[int] = None,
248
+ eos_token_id: Optional[Union[int, List[int]]] = None,
249
+ output_attentions: Optional[bool] = None,
250
+ output_hidden_states: Optional[bool] = None,
251
+ output_scores: Optional[bool] = None,
252
+ return_dict_in_generate: Optional[bool] = None,
253
+ synced_gpus: bool = False,
254
+ streamer: Optional["BaseStreamer"] = None,
255
+ **model_kwargs,
256
+ ) -> Union[GreedySearchOutput, torch.LongTensor]:
257
+
258
+ return self.greedy_search(
259
+ input_ids,
260
+ logits_processor,
261
+ stopping_criteria,
262
+ max_length,
263
+ pad_token_id,
264
+ eos_token_id,
265
+ output_attentions,
266
+ output_hidden_states,
267
+ output_scores,
268
+ return_dict_in_generate,
269
+ synced_gpus,
270
+ streamer,
271
+ **model_kwargs,
272
+ )
273
+ def _beam_search(
274
+ self,
275
+ input_ids: torch.LongTensor,
276
+ beam_scorer: BeamScorer,
277
+ logits_processor: Optional[LogitsProcessorList] = None,
278
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
279
+ max_length: Optional[int] = None,
280
+ pad_token_id: Optional[int] = None,
281
+ eos_token_id: Optional[Union[int, List[int]]] = None,
282
+ output_attentions: Optional[bool] = None,
283
+ output_hidden_states: Optional[bool] = None,
284
+ output_scores: Optional[bool] = None,
285
+ return_dict_in_generate: Optional[bool] = None,
286
+ synced_gpus: bool = False,
287
+ **model_kwargs,
288
+ ) -> Union[BeamSearchOutput, torch.LongTensor]:
289
+
290
+ return self.beam_search(
291
+ input_ids,
292
+ beam_scorer,
293
+ logits_processor,
294
+ stopping_criteria,
295
+ max_length,
296
+ pad_token_id,
297
+ eos_token_id,
298
+ output_attentions,
299
+ output_hidden_states,
300
+ output_scores,
301
+ return_dict_in_generate,
302
+ synced_gpus,
303
+ **model_kwargs,
304
+ )
305
+
306
+ def beam_search(
307
+ self,
308
+ input_ids: torch.LongTensor,
309
+ beam_scorer: BeamScorer,
310
+ logits_processor: Optional[LogitsProcessorList] = None,
311
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
312
+ max_length: Optional[int] = None,
313
+ pad_token_id: Optional[int] = None,
314
+ eos_token_id: Optional[Union[int, List[int]]] = None,
315
+ output_attentions: Optional[bool] = None,
316
+ output_hidden_states: Optional[bool] = None,
317
+ output_scores: Optional[bool] = None,
318
+ return_dict_in_generate: Optional[bool] = None,
319
+ synced_gpus: bool = False,
320
+ **model_kwargs,
321
+ ) -> Union[BeamSearchOutput, torch.LongTensor]:
322
+ '''
323
+ This function is an edited version of the beam_search function from HuggingFace's transformers
324
+ https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py
325
+ '''
326
+ # init values
327
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
328
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
329
+ if max_length is not None:
330
+ warnings.warn(
331
+ "`max_length` is deprecated in this function, use"
332
+ " `stopping_criteria=StoppingCriteriaList(MaxLengthCriteria(max_length=max_length))` instead.",
333
+ UserWarning,
334
+ )
335
+ stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
336
+ if len(stopping_criteria) == 0:
337
+ warnings.warn("You don't have defined any stopping_criteria, this will likely loop forever", UserWarning)
338
+ pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
339
+ eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
340
+ if isinstance(eos_token_id, int):
341
+ eos_token_id = [eos_token_id]
342
+ output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
343
+ output_attentions = (
344
+ output_attentions if output_attentions is not None else self.generation_config.output_attentions
345
+ )
346
+ output_hidden_states = (
347
+ output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
348
+ )
349
+ return_dict_in_generate = (
350
+ return_dict_in_generate
351
+ if return_dict_in_generate is not None
352
+ else self.generation_config.return_dict_in_generate
353
+ )
354
+
355
+ batch_size = len(beam_scorer._beam_hyps)
356
+ num_beams = beam_scorer.num_beams
357
+
358
+ batch_beam_size, cur_len = input_ids.shape
359
+
360
+ if num_beams * batch_size != batch_beam_size:
361
+ raise ValueError(
362
+ f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
363
+ )
364
+
365
+ # init attention / hidden states / scores tuples
366
+ scores = () if (return_dict_in_generate and output_scores) else None
367
+ beam_indices = (
368
+ tuple(() for _ in range(batch_beam_size)) if (return_dict_in_generate and output_scores) else None
369
+ )
370
+ decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
371
+ cross_attentions = () if (return_dict_in_generate and output_attentions) else None
372
+ decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
373
+
374
+ # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
375
+ if return_dict_in_generate and self.config.is_encoder_decoder:
376
+ encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
377
+ encoder_hidden_states = (
378
+ model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
379
+ )
380
+
381
+ # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens
382
+ # of the first beam are considered to avoid sampling the exact same tokens across all beams.
383
+ beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
384
+ beam_scores[:, 1:] = -1e9
385
+ beam_scores = beam_scores.view((batch_size * num_beams,))
386
+
387
+ this_peer_finished = False # used by synced_gpus only
388
+ while True:
389
+ if synced_gpus:
390
+ # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
391
+ # The following logic allows an early break if all peers finished generating their sequence
392
+ this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
393
+ # send 0.0 if we finished, 1.0 otherwise
394
+ dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
395
+ # did all peers finish? the reduced sum will be 0.0 then
396
+ if this_peer_finished_flag.item() == 0.0:
397
+ break
398
+
399
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
400
+
401
+ outputs = self(
402
+ **model_inputs,
403
+ return_dict=True,
404
+ output_attentions=output_attentions,
405
+ output_hidden_states=output_hidden_states,
406
+ )
407
+
408
+ if synced_gpus and this_peer_finished:
409
+ cur_len = cur_len + 1
410
+ continue # don't waste resources running the code we don't need
411
+
412
+ next_token_logits = outputs.logits[:, -1, :]
413
+ # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id`
414
+ # cannot be generated both before and after the `nn.functional.log_softmax` operation.
415
+ # next_token_logits = self.adjust_logits_during_generation(next_token_logits, cur_len=cur_len)
416
+ next_token_scores = nn.functional.log_softmax(
417
+ next_token_logits, dim=-1
418
+ ) # (batch_size * num_beams, vocab_size)
419
+
420
+ next_token_scores_processed = logits_processor(input_ids, next_token_scores)
421
+ # next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(next_token_scores)
422
+ next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(
423
+ next_token_scores_processed
424
+ )
425
+
426
+ # Store scores, attentions and hidden_states when required
427
+ if return_dict_in_generate:
428
+ if output_scores:
429
+ scores += (next_token_scores_processed,)
430
+ if output_attentions:
431
+ decoder_attentions += (
432
+ (outputs.decoder_attentions,) if not self.config.is_encoder_decoder else (outputs.attentions,)
433
+ )
434
+ if self.config.is_encoder_decoder:
435
+ cross_attentions += (outputs.cross_attentions,)
436
+
437
+ if output_hidden_states:
438
+ decoder_hidden_states += (
439
+ (outputs.decoder_hidden_states,)
440
+ if self.config.is_encoder_decoder
441
+ else (outputs.hidden_states,)
442
+ )
443
+
444
+ # reshape for beam search
445
+ vocab_size = next_token_scores.shape[-1]
446
+ next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
447
+
448
+
449
+
450
+ # Sample 2 next tokens for each beam (so we have some spare tokens and match output of beam search)
451
+ next_token_scores, next_tokens = torch.topk(
452
+ next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True
453
+ )
454
+
455
+ next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
456
+ next_tokens = next_tokens % vocab_size
457
+
458
+ # stateless
459
+ beam_outputs = beam_scorer.process(
460
+ input_ids,
461
+ next_token_scores,
462
+ next_tokens,
463
+ next_indices,
464
+ pad_token_id=pad_token_id,
465
+ eos_token_id=eos_token_id,
466
+ beam_indices=beam_indices,
467
+ )
468
+
469
+ beam_scores = beam_outputs["next_beam_scores"]
470
+ beam_next_tokens = beam_outputs["next_beam_tokens"]
471
+ beam_idx = beam_outputs["next_beam_indices"]
472
+
473
+ input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
474
+
475
+ model_kwargs = self._update_model_kwargs_for_generation(
476
+ outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
477
+ )
478
+ if model_kwargs["past_key_values"] is not None:
479
+ model_kwargs["past_key_values"] = self._reorder_cache(model_kwargs["past_key_values"], beam_idx)
480
+
481
+ if return_dict_in_generate and output_scores:
482
+ beam_indices = tuple((beam_indices[beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices))))
483
+
484
+ # increase cur_len
485
+ cur_len = cur_len + 1
486
+
487
+ try:
488
+ if beam_scorer.is_done or stopping_criteria(input_ids, scores):
489
+ if not synced_gpus:
490
+ break
491
+ else:
492
+ this_peer_finished = True
493
+ except:
494
+ if beam_scorer.is_done or all(stopping_criteria(input_ids, scores)):
495
+ if not synced_gpus:
496
+ break
497
+ else:
498
+ this_peer_finished = True
499
+
500
+
501
+ sequence_outputs = beam_scorer.finalize(
502
+ input_ids,
503
+ beam_scores,
504
+ next_tokens,
505
+ next_indices,
506
+ pad_token_id=pad_token_id,
507
+ eos_token_id=eos_token_id,
508
+ max_length=stopping_criteria.max_length,
509
+ beam_indices=beam_indices,
510
+ )
511
+
512
+ if return_dict_in_generate:
513
+ if not output_scores:
514
+ sequence_outputs["sequence_scores"] = None
515
+
516
+ if self.config.is_encoder_decoder:
517
+ return BeamSearchEncoderDecoderOutput(
518
+ sequences=sequence_outputs["sequences"],
519
+ sequences_scores=sequence_outputs["sequence_scores"],
520
+ scores=scores,
521
+ beam_indices=sequence_outputs["beam_indices"],
522
+ encoder_attentions=encoder_attentions,
523
+ encoder_hidden_states=encoder_hidden_states,
524
+ decoder_attentions=decoder_attentions,
525
+ cross_attentions=cross_attentions,
526
+ decoder_hidden_states=decoder_hidden_states,
527
+ )
528
+ else:
529
+ return BeamSearchDecoderOnlyOutput(
530
+ sequences=sequence_outputs["sequences"],
531
+ sequences_scores=sequence_outputs["sequence_scores"],
532
+ scores=scores,
533
+ beam_indices=sequence_outputs["beam_indices"],
534
+ attentions=decoder_attentions,
535
+ hidden_states=decoder_hidden_states,
536
+ )
537
+ else:
538
+ return sequence_outputs["sequences"]
539
+
540
+
541
+ class CABlock(nn.Module):
542
+ '''
543
+ This function is an edited version of the gpt2 decoder block function from HuggingFace's transformers
544
+ https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
545
+ '''
546
+ def __init__(self, config, layer_idx=None):
547
+ super().__init__()
548
+ hidden_size = config.hidden_size
549
+ inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
550
+
551
+ self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
552
+
553
+ self.crossattention = GPT2Attention(config, is_cross_attention=True, layer_idx=layer_idx)
554
+ self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
555
+
556
+ self.mlp = GPT2MLP(inner_dim, config)
557
+
558
+ def forward(
559
+ self,
560
+ hidden_states: Optional[Tuple[torch.FloatTensor]],
561
+ layer_past: Optional[Tuple[torch.Tensor]] = None,
562
+ attention_mask: Optional[torch.FloatTensor] = None,
563
+ head_mask: Optional[torch.FloatTensor] = None,
564
+ encoder_hidden_states: Optional[torch.Tensor] = None,
565
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
566
+ use_cache: Optional[bool] = False,
567
+ output_attentions: Optional[bool] = False,
568
+ ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
569
+
570
+
571
+ residual = hidden_states
572
+ hidden_states = self.ln_cross_attn(hidden_states)
573
+ cross_attn_outputs = self.crossattention(
574
+ hidden_states,
575
+ attention_mask=attention_mask,
576
+ head_mask=head_mask,
577
+ encoder_hidden_states=encoder_hidden_states,
578
+ encoder_attention_mask=encoder_attention_mask,
579
+ output_attentions=output_attentions,
580
+ )
581
+ attn_output = cross_attn_outputs[0]
582
+ # residual connection
583
+ hidden_states = residual + attn_output
584
+
585
+ residual = hidden_states
586
+ hidden_states = self.ln_2(hidden_states)
587
+ feed_forward_hidden_states = self.mlp(hidden_states)
588
+ # residual connection
589
+ hidden_states = residual + feed_forward_hidden_states
590
+
591
+ return (hidden_states,)
592
+
593
+ class Prot2TextTrainer(Seq2SeqTrainer):
594
+ '''
595
+ This function is an edited version of the Seq2SeqTrainer from HuggingFace's transformers
596
+ '''
597
+ def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
598
+ if self.args.world_size > 1:
599
+ eval_sampler = DistributedSampler(self.eval_dataset, num_replicas=self.args.world_size, rank=self.args.process_index)
600
+ else:
601
+ eval_sampler = None
602
+ return DataLoader(
603
+ self.eval_dataset,
604
+ batch_size=self.args.eval_batch_size,
605
+ collate_fn=None,
606
+ num_workers=self.args.dataloader_num_workers,
607
+ pin_memory=self.args.dataloader_pin_memory,
608
+ sampler=eval_sampler,
609
+ )
610
+ def get_train_dataloader(self) -> DataLoader:
611
+ if self.args.world_size > 1:
612
+ train_sampler = DistributedSampler(self.train_dataset, num_replicas=self.args.world_size, rank=self.args.process_index)
613
+ else:
614
+ train_sampler = None
615
+ return DataLoader(
616
+ self.train_dataset,
617
+ batch_size=self.args.per_device_train_batch_size,
618
+ collate_fn=None,
619
+ num_workers=self.args.dataloader_num_workers,
620
+ pin_memory=self.args.dataloader_pin_memory,
621
+ sampler=train_sampler,
622
+ )
623
+ def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
624
+ """
625
+ Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and
626
+ handling potential state.
627
+ """
628
+ inputs = self._prepare_input(inputs)
629
+ if len(inputs) == 0:
630
+ raise ValueError(
631
+ "The batch received was empty, your model won't be able to train on it. Double-check that your "
632
+ f"training dataset contains keys expected by the model: {','.join(self._signature_columns)}."
633
+ )
634
+ if self.args.past_index >= 0 and self._past is not None:
635
+ inputs["mems"] = self._past
636
+
637
+ inputs = inputs.to_dict()
638
+ inputs['edge_type'] = torch.cat([torch.tensor(inputs['edge_type'][i]) for i in range(len(inputs['edge_type']))], dim=0)
639
+ inputs['edge_type'] = torch.argmax(inputs['edge_type'], dim=1)
640
+ inputs = {k: v.to(device=self.args.device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
641
+ return inputs
642
+
643
+ def prediction_step(
644
+ self,
645
+ model: nn.Module,
646
+ inputs: Dict[str, Union[torch.Tensor, Any]],
647
+ prediction_loss_only: bool,
648
+ ignore_keys: Optional[List[str]] = None,
649
+ ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
650
+ """
651
+ Perform an evaluation step on `model` using `inputs`.
652
+
653
+ Subclass and override to inject custom behavior.
654
+
655
+ Args:
656
+ model (`nn.Module`):
657
+ The model to evaluate.
658
+ inputs (`Dict[str, Union[torch.Tensor, Any]]`):
659
+ The inputs and targets of the model.
660
+
661
+ The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
662
+ argument `labels`. Check your model's documentation for all accepted arguments.
663
+ prediction_loss_only (`bool`):
664
+ Whether or not to return the loss only.
665
+
666
+ Return:
667
+ Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and
668
+ labels (each being optional).
669
+ """
670
+
671
+ if not self.args.predict_with_generate or prediction_loss_only:
672
+ return super().prediction_step(
673
+ model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
674
+ )
675
+
676
+ has_labels = "labels" in inputs
677
+ inputs = self._prepare_inputs(inputs)
678
+
679
+ # XXX: adapt synced_gpus for fairscale as well
680
+ gen_kwargs = self._gen_kwargs.copy()
681
+ if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None:
682
+ gen_kwargs["max_length"] = self.model.config.max_length
683
+ gen_kwargs["num_beams"] = (
684
+ gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.model.config.num_beams
685
+ )
686
+ default_synced_gpus = True if is_deepspeed_zero3_enabled() else False
687
+ gen_kwargs["synced_gpus"] = (
688
+ gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus") is not None else default_synced_gpus
689
+ )
690
+
691
+ if "attention_mask" in inputs:
692
+ gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
693
+ if "global_attention_mask" in inputs:
694
+ gen_kwargs["global_attention_mask"] = inputs.get("global_attention_mask", None)
695
+
696
+ generation_inputs = None
697
+ gen_kwargs['x'] = inputs.get('x', None)
698
+ gen_kwargs['edge_index'] = inputs.get('edge_index', None)
699
+ gen_kwargs['edge_type'] = inputs.get('edge_type', None)
700
+ gen_kwargs['batch'] = inputs.get('batch', None)
701
+ gen_kwargs['encoder_input_ids'] = inputs.get('encoder_input_ids', None)
702
+ gen_kwargs['decoder_input_ids'] = inputs.get('decoder_input_ids', None)[:,0:1]
703
+ gen_kwargs["decoder_attention_mask"] = torch.ones(gen_kwargs['decoder_input_ids'].shape[0], 1).to(self.args.device)
704
+
705
+ generated_tokens = self.model.generate(
706
+ generation_inputs,
707
+ **gen_kwargs,
708
+ )
709
+ # in case the batch is shorter than max length, the output should be padded
710
+ if gen_kwargs.get("max_length") is not None and generated_tokens.shape[-1] < gen_kwargs["max_length"]:
711
+ generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_length"])
712
+ elif gen_kwargs.get("max_new_tokens") is not None and generated_tokens.shape[-1] < (
713
+ gen_kwargs["max_new_tokens"] + 1
714
+ ):
715
+ generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_new_tokens"] + 1)
716
+
717
+ with torch.no_grad():
718
+ if has_labels:
719
+ with self.compute_loss_context_manager():
720
+ outputs = model(**inputs)
721
+ if self.label_smoother is not None:
722
+ loss = self.label_smoother(outputs, inputs["labels"]).mean().detach()
723
+ else:
724
+ loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).mean().detach()
725
+ else:
726
+ loss = None
727
+
728
+ if self.args.prediction_loss_only:
729
+ return (loss, None, None)
730
+
731
+ if has_labels:
732
+ labels = inputs["labels"]
733
+ if gen_kwargs.get("max_length") is not None and labels.shape[-1] < gen_kwargs["max_length"]:
734
+ labels = self._pad_tensors_to_max_len(labels, gen_kwargs["max_length"])
735
+ elif gen_kwargs.get("max_new_tokens") is not None and labels.shape[-1] < (
736
+ gen_kwargs["max_new_tokens"] + 1
737
+ ):
738
+ labels = self._pad_tensors_to_max_len(labels, (gen_kwargs["max_new_tokens"] + 1))
739
+ else:
740
+ labels = None
741
+
742
+ return (loss, generated_tokens, labels)
utils_convert.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ from biopandas.pdb import PandasPdb
3
+
4
+ pdb_order = [
5
+ "record_name",
6
+ "atom_number",
7
+ "blank_1",
8
+ "atom_name",
9
+ "alt_loc",
10
+ "residue_name",
11
+ "blank_2",
12
+ "chain_id",
13
+ "residue_number",
14
+ "insertion",
15
+ "blank_3",
16
+ "x_coord",
17
+ "y_coord",
18
+ "z_coord",
19
+ "occupancy",
20
+ "b_factor",
21
+ "blank_4",
22
+ "segment_id",
23
+ "element_symbol",
24
+ "charge",
25
+ "line_idx",
26
+ ]
27
+ mmcif_read = {
28
+ "group_PDB": "record_name",
29
+ "id": "atom_number",
30
+ "auth_atom_id": "atom_name",
31
+ "auth_comp_id": "residue_name",
32
+ "auth_asym_id": "chain_id",
33
+ "auth_seq_id": "residue_number",
34
+ "Cartn_x": "x_coord",
35
+ "Cartn_y": "y_coord",
36
+ "Cartn_z": "z_coord",
37
+ "occupancy": "occupancy",
38
+ "B_iso_or_equiv": "b_factor",
39
+ "type_symbol": "element_symbol",
40
+ }
41
+
42
+ nonefields = [
43
+ "blank_1",
44
+ "alt_loc",
45
+ "blank_2",
46
+ "insertion",
47
+ "blank_3",
48
+ "blank_4",
49
+ "segment_id",
50
+ "charge",
51
+ "line_idx",
52
+ ]
53
+
54
+
55
+ def biopandas_mmcif2pdb(pandasmmcif, model_index = 1):
56
+ """
57
+ Converts the ATOM and HETATM dataframes of PandasMmcif() to PandasPdb() format.
58
+ """
59
+ pandaspdb = PandasPdb()
60
+ for a in ["ATOM", "HETATM"]:
61
+ dfa = pandasmmcif.df[a]
62
+ dfa = dfa.loc[dfa.pdbx_PDB_model_num == model_index]
63
+ if a =='ATOM':
64
+ if len(dfa) == 0:
65
+ raise ValueError(f"No model found for index: {model_index}")
66
+ # keep only those fields found in pdb
67
+ dfa = dfa[mmcif_read.keys()]
68
+ # rename fields
69
+ dfa = dfa.rename(columns=mmcif_read)
70
+ # add empty fields
71
+ for i in nonefields:
72
+ dfa[i] = ""
73
+ dfa["charge"] = np.nan
74
+ # reorder columns to PandasPdb order
75
+ dfa = dfa[pdb_order]
76
+ pandaspdb.df[a] = dfa
77
+
78
+ # update line_idx
79
+ pandaspdb.df["ATOM"]["line_idx"] = pandaspdb.df["ATOM"].index.values
80
+ pandaspdb.df["HETATM"]["line_idx"] = pandaspdb.df["HETATM"].index
81
+
82
+ return pandaspdb
utils_dataset.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import csv
3
+
4
+ def load_GO_annot(filename):
5
+ # Load GO annotations
6
+ onts = ['mf', 'bp', 'cc']
7
+ prot2annot = {}
8
+ goterms = {ont: [] for ont in onts}
9
+ gonames = {ont: [] for ont in onts}
10
+ with open(filename, mode='r') as tsvfile:
11
+ reader = csv.reader(tsvfile, delimiter='\t')
12
+
13
+ # molecular function
14
+ next(reader, None) # skip the headers
15
+ goterms[onts[0]] = next(reader)
16
+ next(reader, None) # skip the headers
17
+ gonames[onts[0]] = next(reader)
18
+
19
+ # biological process
20
+ next(reader, None) # skip the headers
21
+ goterms[onts[1]] = next(reader)
22
+ next(reader, None) # skip the headers
23
+ gonames[onts[1]] = next(reader)
24
+
25
+ # cellular component
26
+ next(reader, None) # skip the headers
27
+ goterms[onts[2]] = next(reader)
28
+ next(reader, None) # skip the headers
29
+ gonames[onts[2]] = next(reader)
30
+
31
+ next(reader, None) # skip the headers
32
+ counts = {ont: np.zeros(len(goterms[ont]), dtype=float) for ont in onts}
33
+ for row in reader:
34
+ prot, prot_goterms = row[0], row[1:]
35
+ prot2annot[prot] = {ont: [] for ont in onts}
36
+ for i in range(3):
37
+ goterm_indices = [goterms[onts[i]].index(goterm) for goterm in prot_goterms[i].split(',') if goterm != '']
38
+ prot2annot[prot][onts[i]] = np.zeros(len(goterms[onts[i]]))
39
+ prot2annot[prot][onts[i]][goterm_indices] = 1.0
40
+ counts[onts[i]][goterm_indices] += 1.0
41
+ return prot2annot, goterms, gonames, counts
42
+
43
+
44
+ def load_EC_annot(filename):
45
+ # Load EC annotations """
46
+ prot2annot = {}
47
+ with open(filename, mode='r') as tsvfile:
48
+ reader = csv.reader(tsvfile, delimiter='\t')
49
+
50
+ # molecular function
51
+ next(reader, None) # skip the headers
52
+ ec_numbers = {'ec': next(reader)}
53
+ next(reader, None) # skip the headers
54
+ counts = {'ec': np.zeros(len(ec_numbers['ec']), dtype=float)}
55
+ for row in reader:
56
+ prot, prot_ec_numbers = row[0], row[1]
57
+ ec_indices = [ec_numbers['ec'].index(ec_num) for ec_num in prot_ec_numbers.split(',')]
58
+ prot2annot[prot] = {'ec': np.zeros(len(ec_numbers['ec']), dtype=np.int64)}
59
+ prot2annot[prot]['ec'][ec_indices] = 1.0
60
+ counts['ec'][ec_indices] += 1