We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Data missing for JTNN
## ❓ Questions and Help

   Hello,I come across the problem that the website you put your jtnn.zip is missing.Can you offer me a new website?
  thanks!


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/pytorch/jtnn/jtnn/datautils.py`
Content:
```
1 import torch
2 from torch.utils.data import Dataset
3 import numpy as np
4 
5 import dgl
6 from dgl.data.utils import download, extract_archive, get_download_dir
7 from .mol_tree_nx import DGLMolTree
8 from .mol_tree import Vocab
9 
10 from .mpn import mol2dgl_single as mol2dgl_enc
11 from .jtmpn import mol2dgl_single as mol2dgl_dec
12 from .jtmpn import ATOM_FDIM as ATOM_FDIM_DEC
13 from .jtmpn import BOND_FDIM as BOND_FDIM_DEC
14 
15 _url = 'https://www.dropbox.com/s/4ypr0e0abcbsvoh/jtnn.zip?dl=1'
16 
17 def _unpack_field(examples, field):
18     return [e[field] for e in examples]
19 
20 def _set_node_id(mol_tree, vocab):
21     wid = []
22     for i, node in enumerate(mol_tree.nodes_dict):
23         mol_tree.nodes_dict[node]['idx'] = i
24         wid.append(vocab.get_index(mol_tree.nodes_dict[node]['smiles']))
25 
26     return wid
27 
28 class JTNNDataset(Dataset):
29     def __init__(self, data, vocab, training=True):
30         self.dir = get_download_dir()
31         self.zip_file_path='{}/jtnn.zip'.format(self.dir)
32         download(_url, path=self.zip_file_path)
33         extract_archive(self.zip_file_path, '{}/jtnn'.format(self.dir))
34         print('Loading data...')
35         data_file = '{}/jtnn/{}.txt'.format(self.dir, data)
36         with open(data_file) as f:
37             self.data = [line.strip("\r\n ").split()[0] for line in f]
38         self.vocab_file = '{}/jtnn/{}.txt'.format(self.dir, vocab)
39         print('Loading finished.')
40         print('\tNum samples:', len(self.data))
41         print('\tVocab file:', self.vocab_file)
42         self.training = training
43         self.vocab = Vocab([x.strip("\r\n ") for x in open(self.vocab_file)])
44 
45     def __len__(self):
46         return len(self.data)
47 
48     def __getitem__(self, idx):
49         smiles = self.data[idx]
50         mol_tree = DGLMolTree(smiles)
51         mol_tree.recover()
52         mol_tree.assemble()
53 
54         wid = _set_node_id(mol_tree, self.vocab)
55 
56         # prebuild the molecule graph
57         mol_graph, atom_x_enc, bond_x_enc = mol2dgl_enc(mol_tree.smiles)
58 
59         result = {
60                 'mol_tree': mol_tree,
61                 'mol_graph': mol_graph,
62                 'atom_x_enc': atom_x_enc,
63                 'bond_x_enc': bond_x_enc,
64                 'wid': wid,
65                 }
66 
67         if not self.training:
68             return result
69 
70         # prebuild the candidate graph list
71         cands = []
72         for node_id, node in mol_tree.nodes_dict.items():
73             # fill in ground truth
74             if node['label'] not in node['cands']:
75                 node['cands'].append(node['label'])
76                 node['cand_mols'].append(node['label_mol'])
77 
78             if node['is_leaf'] or len(node['cands']) == 1:
79                 continue
80             cands.extend([(cand, mol_tree, node_id)
81                          for cand in node['cand_mols']])
82         if len(cands) > 0:
83             cand_graphs, atom_x_dec, bond_x_dec, tree_mess_src_e, \
84                     tree_mess_tgt_e, tree_mess_tgt_n = mol2dgl_dec(cands)
85         else:
86             cand_graphs = []
87             atom_x_dec = torch.zeros(0, ATOM_FDIM_DEC)
88             bond_x_dec = torch.zeros(0, BOND_FDIM_DEC)
89             tree_mess_src_e = torch.zeros(0, 2).long()
90             tree_mess_tgt_e = torch.zeros(0, 2).long()
91             tree_mess_tgt_n = torch.zeros(0).long()
92 
93         # prebuild the stereoisomers
94         cands = mol_tree.stereo_cands
95         if len(cands) > 1:
96             if mol_tree.smiles3D not in cands:
97                 cands.append(mol_tree.smiles3D)
98 
99             stereo_graphs = [mol2dgl_enc(c) for c in cands]
100             stereo_cand_graphs, stereo_atom_x_enc, stereo_bond_x_enc = \
101                     zip(*stereo_graphs)
102             stereo_atom_x_enc = torch.cat(stereo_atom_x_enc)
103             stereo_bond_x_enc = torch.cat(stereo_bond_x_enc)
104             stereo_cand_label = [(cands.index(mol_tree.smiles3D), len(cands))]
105         else:
106             stereo_cand_graphs = []
107             stereo_atom_x_enc = torch.zeros(0, atom_x_enc.shape[1])
108             stereo_bond_x_enc = torch.zeros(0, bond_x_enc.shape[1])
109             stereo_cand_label = []
110 
111         result.update({
112             'cand_graphs': cand_graphs,
113             'atom_x_dec': atom_x_dec,
114             'bond_x_dec': bond_x_dec,
115             'tree_mess_src_e': tree_mess_src_e,
116             'tree_mess_tgt_e': tree_mess_tgt_e,
117             'tree_mess_tgt_n': tree_mess_tgt_n,
118             'stereo_cand_graphs': stereo_cand_graphs,
119             'stereo_atom_x_enc': stereo_atom_x_enc,
120             'stereo_bond_x_enc': stereo_bond_x_enc,
121             'stereo_cand_label': stereo_cand_label,
122             })
123 
124         return result
125 
126 class JTNNCollator(object):
127     def __init__(self, vocab, training):
128         self.vocab = vocab
129         self.training = training
130 
131     @staticmethod
132     def _batch_and_set(graphs, atom_x, bond_x, flatten):
133         if flatten:
134             graphs = [g for f in graphs for g in f]
135         graph_batch = dgl.batch(graphs)
136         graph_batch.ndata['x'] = atom_x
137         graph_batch.edata.update({
138             'x': bond_x,
139             'src_x': atom_x.new(bond_x.shape[0], atom_x.shape[1]).zero_(),
140             })
141         return graph_batch
142 
143     def __call__(self, examples):
144         # get list of trees
145         mol_trees = _unpack_field(examples, 'mol_tree')
146         wid = _unpack_field(examples, 'wid')
147         for _wid, mol_tree in zip(wid, mol_trees):
148             mol_tree.ndata['wid'] = torch.LongTensor(_wid)
149 
150         # TODO: either support pickling or get around ctypes pointers using scipy
151         # batch molecule graphs
152         mol_graphs = _unpack_field(examples, 'mol_graph')
153         atom_x = torch.cat(_unpack_field(examples, 'atom_x_enc'))
154         bond_x = torch.cat(_unpack_field(examples, 'bond_x_enc'))
155         mol_graph_batch = self._batch_and_set(mol_graphs, atom_x, bond_x, False)
156 
157         result = {
158                 'mol_trees': mol_trees,
159                 'mol_graph_batch': mol_graph_batch,
160                 }
161 
162         if not self.training:
163             return result
164 
165         # batch candidate graphs
166         cand_graphs = _unpack_field(examples, 'cand_graphs')
167         cand_batch_idx = []
168         atom_x = torch.cat(_unpack_field(examples, 'atom_x_dec'))
169         bond_x = torch.cat(_unpack_field(examples, 'bond_x_dec'))
170         tree_mess_src_e = _unpack_field(examples, 'tree_mess_src_e')
171         tree_mess_tgt_e = _unpack_field(examples, 'tree_mess_tgt_e')
172         tree_mess_tgt_n = _unpack_field(examples, 'tree_mess_tgt_n')
173 
174         n_graph_nodes = 0
175         n_tree_nodes = 0
176         for i in range(len(cand_graphs)):
177             tree_mess_tgt_e[i] += n_graph_nodes
178             tree_mess_src_e[i] += n_tree_nodes
179             tree_mess_tgt_n[i] += n_graph_nodes
180             n_graph_nodes += sum(g.number_of_nodes() for g in cand_graphs[i])
181             n_tree_nodes += mol_trees[i].number_of_nodes()
182             cand_batch_idx.extend([i] * len(cand_graphs[i]))
183         tree_mess_tgt_e = torch.cat(tree_mess_tgt_e)
184         tree_mess_src_e = torch.cat(tree_mess_src_e)
185         tree_mess_tgt_n = torch.cat(tree_mess_tgt_n)
186 
187         cand_graph_batch = self._batch_and_set(cand_graphs, atom_x, bond_x, True)
188 
189         # batch stereoisomers
190         stereo_cand_graphs = _unpack_field(examples, 'stereo_cand_graphs')
191         atom_x = torch.cat(_unpack_field(examples, 'stereo_atom_x_enc'))
192         bond_x = torch.cat(_unpack_field(examples, 'stereo_bond_x_enc'))
193         stereo_cand_batch_idx = []
194         for i in range(len(stereo_cand_graphs)):
195             stereo_cand_batch_idx.extend([i] * len(stereo_cand_graphs[i]))
196 
197         if len(stereo_cand_batch_idx) > 0:
198             stereo_cand_labels = [
199                     (label, length)
200                     for ex in _unpack_field(examples, 'stereo_cand_label')
201                     for label, length in ex
202                     ]
203             stereo_cand_labels, stereo_cand_lengths = zip(*stereo_cand_labels)
204             stereo_cand_graph_batch = self._batch_and_set(
205                     stereo_cand_graphs, atom_x, bond_x, True)
206         else:
207             stereo_cand_labels = []
208             stereo_cand_lengths = []
209             stereo_cand_graph_batch = None
210             stereo_cand_batch_idx = []
211 
212         result.update({
213             'cand_graph_batch': cand_graph_batch,
214             'cand_batch_idx': cand_batch_idx,
215             'tree_mess_tgt_e': tree_mess_tgt_e,
216             'tree_mess_src_e': tree_mess_src_e,
217             'tree_mess_tgt_n': tree_mess_tgt_n,
218             'stereo_cand_graph_batch': stereo_cand_graph_batch,
219             'stereo_cand_batch_idx': stereo_cand_batch_idx,
220             'stereo_cand_labels': stereo_cand_labels,
221             'stereo_cand_lengths': stereo_cand_lengths,
222             })
223 
224         return result
225
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/pytorch/jtnn/jtnn/datautils.py b/examples/pytorch/jtnn/jtnn/datautils.py
--- a/examples/pytorch/jtnn/jtnn/datautils.py
+++ b/examples/pytorch/jtnn/jtnn/datautils.py
@@ -12,7 +12,7 @@
 from .jtmpn import ATOM_FDIM as ATOM_FDIM_DEC
 from .jtmpn import BOND_FDIM as BOND_FDIM_DEC
 
-_url = 'https://www.dropbox.com/s/4ypr0e0abcbsvoh/jtnn.zip?dl=1'
+_url = 'https://s3-ap-southeast-1.amazonaws.com/dgl-data-cn/dataset/jtnn.zip'
 
 def _unpack_field(examples, field):
     return [e[field] for e in examples]

{"golden_diff": "diff --git a/examples/pytorch/jtnn/jtnn/datautils.py b/examples/pytorch/jtnn/jtnn/datautils.py\n--- a/examples/pytorch/jtnn/jtnn/datautils.py\n+++ b/examples/pytorch/jtnn/jtnn/datautils.py\n@@ -12,7 +12,7 @@\n from .jtmpn import ATOM_FDIM as ATOM_FDIM_DEC\n from .jtmpn import BOND_FDIM as BOND_FDIM_DEC\n \n-_url = 'https://www.dropbox.com/s/4ypr0e0abcbsvoh/jtnn.zip?dl=1'\n+_url = 'https://s3-ap-southeast-1.amazonaws.com/dgl-data-cn/dataset/jtnn.zip'\n \n def _unpack_field(examples, field):\n     return [e[field] for e in examples]\n", "issue": "Data missing for JTNN\n## \u2753 Questions and Help\r\n\r\n   Hello,I come across the problem that the website you put your jtnn.zip is missing.Can you offer me a new website?\r\n  thanks!\r\n\n", "before_files": [{"content": "import torch\nfrom torch.utils.data import Dataset\nimport numpy as np\n\nimport dgl\nfrom dgl.data.utils import download, extract_archive, get_download_dir\nfrom .mol_tree_nx import DGLMolTree\nfrom .mol_tree import Vocab\n\nfrom .mpn import mol2dgl_single as mol2dgl_enc\nfrom .jtmpn import mol2dgl_single as mol2dgl_dec\nfrom .jtmpn import ATOM_FDIM as ATOM_FDIM_DEC\nfrom .jtmpn import BOND_FDIM as BOND_FDIM_DEC\n\n_url = 'https://www.dropbox.com/s/4ypr0e0abcbsvoh/jtnn.zip?dl=1'\n\ndef _unpack_field(examples, field):\n    return [e[field] for e in examples]\n\ndef _set_node_id(mol_tree, vocab):\n    wid = []\n    for i, node in enumerate(mol_tree.nodes_dict):\n        mol_tree.nodes_dict[node]['idx'] = i\n        wid.append(vocab.get_index(mol_tree.nodes_dict[node]['smiles']))\n\n    return wid\n\nclass JTNNDataset(Dataset):\n    def __init__(self, data, vocab, training=True):\n        self.dir = get_download_dir()\n        self.zip_file_path='{}/jtnn.zip'.format(self.dir)\n        download(_url, path=self.zip_file_path)\n        extract_archive(self.zip_file_path, '{}/jtnn'.format(self.dir))\n        print('Loading data...')\n        data_file = '{}/jtnn/{}.txt'.format(self.dir, data)\n        with open(data_file) as f:\n            self.data = [line.strip(\"\\r\\n \").split()[0] for line in f]\n        self.vocab_file = '{}/jtnn/{}.txt'.format(self.dir, vocab)\n        print('Loading finished.')\n        print('\\tNum samples:', len(self.data))\n        print('\\tVocab file:', self.vocab_file)\n        self.training = training\n        self.vocab = Vocab([x.strip(\"\\r\\n \") for x in open(self.vocab_file)])\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        smiles = self.data[idx]\n        mol_tree = DGLMolTree(smiles)\n        mol_tree.recover()\n        mol_tree.assemble()\n\n        wid = _set_node_id(mol_tree, self.vocab)\n\n        # prebuild the molecule graph\n        mol_graph, atom_x_enc, bond_x_enc = mol2dgl_enc(mol_tree.smiles)\n\n        result = {\n                'mol_tree': mol_tree,\n                'mol_graph': mol_graph,\n                'atom_x_enc': atom_x_enc,\n                'bond_x_enc': bond_x_enc,\n                'wid': wid,\n                }\n\n        if not self.training:\n            return result\n\n        # prebuild the candidate graph list\n        cands = []\n        for node_id, node in mol_tree.nodes_dict.items():\n            # fill in ground truth\n            if node['label'] not in node['cands']:\n                node['cands'].append(node['label'])\n                node['cand_mols'].append(node['label_mol'])\n\n            if node['is_leaf'] or len(node['cands']) == 1:\n                continue\n            cands.extend([(cand, mol_tree, node_id)\n                         for cand in node['cand_mols']])\n        if len(cands) > 0:\n            cand_graphs, atom_x_dec, bond_x_dec, tree_mess_src_e, \\\n                    tree_mess_tgt_e, tree_mess_tgt_n = mol2dgl_dec(cands)\n        else:\n            cand_graphs = []\n            atom_x_dec = torch.zeros(0, ATOM_FDIM_DEC)\n            bond_x_dec = torch.zeros(0, BOND_FDIM_DEC)\n            tree_mess_src_e = torch.zeros(0, 2).long()\n            tree_mess_tgt_e = torch.zeros(0, 2).long()\n            tree_mess_tgt_n = torch.zeros(0).long()\n\n        # prebuild the stereoisomers\n        cands = mol_tree.stereo_cands\n        if len(cands) > 1:\n            if mol_tree.smiles3D not in cands:\n                cands.append(mol_tree.smiles3D)\n\n            stereo_graphs = [mol2dgl_enc(c) for c in cands]\n            stereo_cand_graphs, stereo_atom_x_enc, stereo_bond_x_enc = \\\n                    zip(*stereo_graphs)\n            stereo_atom_x_enc = torch.cat(stereo_atom_x_enc)\n            stereo_bond_x_enc = torch.cat(stereo_bond_x_enc)\n            stereo_cand_label = [(cands.index(mol_tree.smiles3D), len(cands))]\n        else:\n            stereo_cand_graphs = []\n            stereo_atom_x_enc = torch.zeros(0, atom_x_enc.shape[1])\n            stereo_bond_x_enc = torch.zeros(0, bond_x_enc.shape[1])\n            stereo_cand_label = []\n\n        result.update({\n            'cand_graphs': cand_graphs,\n            'atom_x_dec': atom_x_dec,\n            'bond_x_dec': bond_x_dec,\n            'tree_mess_src_e': tree_mess_src_e,\n            'tree_mess_tgt_e': tree_mess_tgt_e,\n            'tree_mess_tgt_n': tree_mess_tgt_n,\n            'stereo_cand_graphs': stereo_cand_graphs,\n            'stereo_atom_x_enc': stereo_atom_x_enc,\n            'stereo_bond_x_enc': stereo_bond_x_enc,\n            'stereo_cand_label': stereo_cand_label,\n            })\n\n        return result\n\nclass JTNNCollator(object):\n    def __init__(self, vocab, training):\n        self.vocab = vocab\n        self.training = training\n\n    @staticmethod\n    def _batch_and_set(graphs, atom_x, bond_x, flatten):\n        if flatten:\n            graphs = [g for f in graphs for g in f]\n        graph_batch = dgl.batch(graphs)\n        graph_batch.ndata['x'] = atom_x\n        graph_batch.edata.update({\n            'x': bond_x,\n            'src_x': atom_x.new(bond_x.shape[0], atom_x.shape[1]).zero_(),\n            })\n        return graph_batch\n\n    def __call__(self, examples):\n        # get list of trees\n        mol_trees = _unpack_field(examples, 'mol_tree')\n        wid = _unpack_field(examples, 'wid')\n        for _wid, mol_tree in zip(wid, mol_trees):\n            mol_tree.ndata['wid'] = torch.LongTensor(_wid)\n\n        # TODO: either support pickling or get around ctypes pointers using scipy\n        # batch molecule graphs\n        mol_graphs = _unpack_field(examples, 'mol_graph')\n        atom_x = torch.cat(_unpack_field(examples, 'atom_x_enc'))\n        bond_x = torch.cat(_unpack_field(examples, 'bond_x_enc'))\n        mol_graph_batch = self._batch_and_set(mol_graphs, atom_x, bond_x, False)\n\n        result = {\n                'mol_trees': mol_trees,\n                'mol_graph_batch': mol_graph_batch,\n                }\n\n        if not self.training:\n            return result\n\n        # batch candidate graphs\n        cand_graphs = _unpack_field(examples, 'cand_graphs')\n        cand_batch_idx = []\n        atom_x = torch.cat(_unpack_field(examples, 'atom_x_dec'))\n        bond_x = torch.cat(_unpack_field(examples, 'bond_x_dec'))\n        tree_mess_src_e = _unpack_field(examples, 'tree_mess_src_e')\n        tree_mess_tgt_e = _unpack_field(examples, 'tree_mess_tgt_e')\n        tree_mess_tgt_n = _unpack_field(examples, 'tree_mess_tgt_n')\n\n        n_graph_nodes = 0\n        n_tree_nodes = 0\n        for i in range(len(cand_graphs)):\n            tree_mess_tgt_e[i] += n_graph_nodes\n            tree_mess_src_e[i] += n_tree_nodes\n            tree_mess_tgt_n[i] += n_graph_nodes\n            n_graph_nodes += sum(g.number_of_nodes() for g in cand_graphs[i])\n            n_tree_nodes += mol_trees[i].number_of_nodes()\n            cand_batch_idx.extend([i] * len(cand_graphs[i]))\n        tree_mess_tgt_e = torch.cat(tree_mess_tgt_e)\n        tree_mess_src_e = torch.cat(tree_mess_src_e)\n        tree_mess_tgt_n = torch.cat(tree_mess_tgt_n)\n\n        cand_graph_batch = self._batch_and_set(cand_graphs, atom_x, bond_x, True)\n\n        # batch stereoisomers\n        stereo_cand_graphs = _unpack_field(examples, 'stereo_cand_graphs')\n        atom_x = torch.cat(_unpack_field(examples, 'stereo_atom_x_enc'))\n        bond_x = torch.cat(_unpack_field(examples, 'stereo_bond_x_enc'))\n        stereo_cand_batch_idx = []\n        for i in range(len(stereo_cand_graphs)):\n            stereo_cand_batch_idx.extend([i] * len(stereo_cand_graphs[i]))\n\n        if len(stereo_cand_batch_idx) > 0:\n            stereo_cand_labels = [\n                    (label, length)\n                    for ex in _unpack_field(examples, 'stereo_cand_label')\n                    for label, length in ex\n                    ]\n            stereo_cand_labels, stereo_cand_lengths = zip(*stereo_cand_labels)\n            stereo_cand_graph_batch = self._batch_and_set(\n                    stereo_cand_graphs, atom_x, bond_x, True)\n        else:\n            stereo_cand_labels = []\n            stereo_cand_lengths = []\n            stereo_cand_graph_batch = None\n            stereo_cand_batch_idx = []\n\n        result.update({\n            'cand_graph_batch': cand_graph_batch,\n            'cand_batch_idx': cand_batch_idx,\n            'tree_mess_tgt_e': tree_mess_tgt_e,\n            'tree_mess_src_e': tree_mess_src_e,\n            'tree_mess_tgt_n': tree_mess_tgt_n,\n            'stereo_cand_graph_batch': stereo_cand_graph_batch,\n            'stereo_cand_batch_idx': stereo_cand_batch_idx,\n            'stereo_cand_labels': stereo_cand_labels,\n            'stereo_cand_lengths': stereo_cand_lengths,\n            })\n\n        return result\n", "path": "examples/pytorch/jtnn/jtnn/datautils.py"}], "after_files": [{"content": "import torch\nfrom torch.utils.data import Dataset\nimport numpy as np\n\nimport dgl\nfrom dgl.data.utils import download, extract_archive, get_download_dir\nfrom .mol_tree_nx import DGLMolTree\nfrom .mol_tree import Vocab\n\nfrom .mpn import mol2dgl_single as mol2dgl_enc\nfrom .jtmpn import mol2dgl_single as mol2dgl_dec\nfrom .jtmpn import ATOM_FDIM as ATOM_FDIM_DEC\nfrom .jtmpn import BOND_FDIM as BOND_FDIM_DEC\n\n_url = 'https://s3-ap-southeast-1.amazonaws.com/dgl-data-cn/dataset/jtnn.zip'\n\ndef _unpack_field(examples, field):\n    return [e[field] for e in examples]\n\ndef _set_node_id(mol_tree, vocab):\n    wid = []\n    for i, node in enumerate(mol_tree.nodes_dict):\n        mol_tree.nodes_dict[node]['idx'] = i\n        wid.append(vocab.get_index(mol_tree.nodes_dict[node]['smiles']))\n\n    return wid\n\nclass JTNNDataset(Dataset):\n    def __init__(self, data, vocab, training=True):\n        self.dir = get_download_dir()\n        self.zip_file_path='{}/jtnn.zip'.format(self.dir)\n        download(_url, path=self.zip_file_path)\n        extract_archive(self.zip_file_path, '{}/jtnn'.format(self.dir))\n        print('Loading data...')\n        data_file = '{}/jtnn/{}.txt'.format(self.dir, data)\n        with open(data_file) as f:\n            self.data = [line.strip(\"\\r\\n \").split()[0] for line in f]\n        self.vocab_file = '{}/jtnn/{}.txt'.format(self.dir, vocab)\n        print('Loading finished.')\n        print('\\tNum samples:', len(self.data))\n        print('\\tVocab file:', self.vocab_file)\n        self.training = training\n        self.vocab = Vocab([x.strip(\"\\r\\n \") for x in open(self.vocab_file)])\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        smiles = self.data[idx]\n        mol_tree = DGLMolTree(smiles)\n        mol_tree.recover()\n        mol_tree.assemble()\n\n        wid = _set_node_id(mol_tree, self.vocab)\n\n        # prebuild the molecule graph\n        mol_graph, atom_x_enc, bond_x_enc = mol2dgl_enc(mol_tree.smiles)\n\n        result = {\n                'mol_tree': mol_tree,\n                'mol_graph': mol_graph,\n                'atom_x_enc': atom_x_enc,\n                'bond_x_enc': bond_x_enc,\n                'wid': wid,\n                }\n\n        if not self.training:\n            return result\n\n        # prebuild the candidate graph list\n        cands = []\n        for node_id, node in mol_tree.nodes_dict.items():\n            # fill in ground truth\n            if node['label'] not in node['cands']:\n                node['cands'].append(node['label'])\n                node['cand_mols'].append(node['label_mol'])\n\n            if node['is_leaf'] or len(node['cands']) == 1:\n                continue\n            cands.extend([(cand, mol_tree, node_id)\n                         for cand in node['cand_mols']])\n        if len(cands) > 0:\n            cand_graphs, atom_x_dec, bond_x_dec, tree_mess_src_e, \\\n                    tree_mess_tgt_e, tree_mess_tgt_n = mol2dgl_dec(cands)\n        else:\n            cand_graphs = []\n            atom_x_dec = torch.zeros(0, ATOM_FDIM_DEC)\n            bond_x_dec = torch.zeros(0, BOND_FDIM_DEC)\n            tree_mess_src_e = torch.zeros(0, 2).long()\n            tree_mess_tgt_e = torch.zeros(0, 2).long()\n            tree_mess_tgt_n = torch.zeros(0).long()\n\n        # prebuild the stereoisomers\n        cands = mol_tree.stereo_cands\n        if len(cands) > 1:\n            if mol_tree.smiles3D not in cands:\n                cands.append(mol_tree.smiles3D)\n\n            stereo_graphs = [mol2dgl_enc(c) for c in cands]\n            stereo_cand_graphs, stereo_atom_x_enc, stereo_bond_x_enc = \\\n                    zip(*stereo_graphs)\n            stereo_atom_x_enc = torch.cat(stereo_atom_x_enc)\n            stereo_bond_x_enc = torch.cat(stereo_bond_x_enc)\n            stereo_cand_label = [(cands.index(mol_tree.smiles3D), len(cands))]\n        else:\n            stereo_cand_graphs = []\n            stereo_atom_x_enc = torch.zeros(0, atom_x_enc.shape[1])\n            stereo_bond_x_enc = torch.zeros(0, bond_x_enc.shape[1])\n            stereo_cand_label = []\n\n        result.update({\n            'cand_graphs': cand_graphs,\n            'atom_x_dec': atom_x_dec,\n            'bond_x_dec': bond_x_dec,\n            'tree_mess_src_e': tree_mess_src_e,\n            'tree_mess_tgt_e': tree_mess_tgt_e,\n            'tree_mess_tgt_n': tree_mess_tgt_n,\n            'stereo_cand_graphs': stereo_cand_graphs,\n            'stereo_atom_x_enc': stereo_atom_x_enc,\n            'stereo_bond_x_enc': stereo_bond_x_enc,\n            'stereo_cand_label': stereo_cand_label,\n            })\n\n        return result\n\nclass JTNNCollator(object):\n    def __init__(self, vocab, training):\n        self.vocab = vocab\n        self.training = training\n\n    @staticmethod\n    def _batch_and_set(graphs, atom_x, bond_x, flatten):\n        if flatten:\n            graphs = [g for f in graphs for g in f]\n        graph_batch = dgl.batch(graphs)\n        graph_batch.ndata['x'] = atom_x\n        graph_batch.edata.update({\n            'x': bond_x,\n            'src_x': atom_x.new(bond_x.shape[0], atom_x.shape[1]).zero_(),\n            })\n        return graph_batch\n\n    def __call__(self, examples):\n        # get list of trees\n        mol_trees = _unpack_field(examples, 'mol_tree')\n        wid = _unpack_field(examples, 'wid')\n        for _wid, mol_tree in zip(wid, mol_trees):\n            mol_tree.ndata['wid'] = torch.LongTensor(_wid)\n\n        # TODO: either support pickling or get around ctypes pointers using scipy\n        # batch molecule graphs\n        mol_graphs = _unpack_field(examples, 'mol_graph')\n        atom_x = torch.cat(_unpack_field(examples, 'atom_x_enc'))\n        bond_x = torch.cat(_unpack_field(examples, 'bond_x_enc'))\n        mol_graph_batch = self._batch_and_set(mol_graphs, atom_x, bond_x, False)\n\n        result = {\n                'mol_trees': mol_trees,\n                'mol_graph_batch': mol_graph_batch,\n                }\n\n        if not self.training:\n            return result\n\n        # batch candidate graphs\n        cand_graphs = _unpack_field(examples, 'cand_graphs')\n        cand_batch_idx = []\n        atom_x = torch.cat(_unpack_field(examples, 'atom_x_dec'))\n        bond_x = torch.cat(_unpack_field(examples, 'bond_x_dec'))\n        tree_mess_src_e = _unpack_field(examples, 'tree_mess_src_e')\n        tree_mess_tgt_e = _unpack_field(examples, 'tree_mess_tgt_e')\n        tree_mess_tgt_n = _unpack_field(examples, 'tree_mess_tgt_n')\n\n        n_graph_nodes = 0\n        n_tree_nodes = 0\n        for i in range(len(cand_graphs)):\n            tree_mess_tgt_e[i] += n_graph_nodes\n            tree_mess_src_e[i] += n_tree_nodes\n            tree_mess_tgt_n[i] += n_graph_nodes\n            n_graph_nodes += sum(g.number_of_nodes() for g in cand_graphs[i])\n            n_tree_nodes += mol_trees[i].number_of_nodes()\n            cand_batch_idx.extend([i] * len(cand_graphs[i]))\n        tree_mess_tgt_e = torch.cat(tree_mess_tgt_e)\n        tree_mess_src_e = torch.cat(tree_mess_src_e)\n        tree_mess_tgt_n = torch.cat(tree_mess_tgt_n)\n\n        cand_graph_batch = self._batch_and_set(cand_graphs, atom_x, bond_x, True)\n\n        # batch stereoisomers\n        stereo_cand_graphs = _unpack_field(examples, 'stereo_cand_graphs')\n        atom_x = torch.cat(_unpack_field(examples, 'stereo_atom_x_enc'))\n        bond_x = torch.cat(_unpack_field(examples, 'stereo_bond_x_enc'))\n        stereo_cand_batch_idx = []\n        for i in range(len(stereo_cand_graphs)):\n            stereo_cand_batch_idx.extend([i] * len(stereo_cand_graphs[i]))\n\n        if len(stereo_cand_batch_idx) > 0:\n            stereo_cand_labels = [\n                    (label, length)\n                    for ex in _unpack_field(examples, 'stereo_cand_label')\n                    for label, length in ex\n                    ]\n            stereo_cand_labels, stereo_cand_lengths = zip(*stereo_cand_labels)\n            stereo_cand_graph_batch = self._batch_and_set(\n                    stereo_cand_graphs, atom_x, bond_x, True)\n        else:\n            stereo_cand_labels = []\n            stereo_cand_lengths = []\n            stereo_cand_graph_batch = None\n            stereo_cand_batch_idx = []\n\n        result.update({\n            'cand_graph_batch': cand_graph_batch,\n            'cand_batch_idx': cand_batch_idx,\n            'tree_mess_tgt_e': tree_mess_tgt_e,\n            'tree_mess_src_e': tree_mess_src_e,\n            'tree_mess_tgt_n': tree_mess_tgt_n,\n            'stereo_cand_graph_batch': stereo_cand_graph_batch,\n            'stereo_cand_batch_idx': stereo_cand_batch_idx,\n            'stereo_cand_labels': stereo_cand_labels,\n            'stereo_cand_lengths': stereo_cand_lengths,\n            })\n\n        return result\n", "path": "examples/pytorch/jtnn/jtnn/datautils.py"}]}