Change max batch_size

app.py

@@ -10,7 +10,11 @@ import output
 
 from rdkit import Chem
 from src import const
-from src.datasets import get_dataloader, collate_with_fragment_edges, parse_molecule, MOADDataset
+from src.datasets import (
+    get_dataloader, collate_with_fragment_edges,
+    collate_with_fragment_without_pocket_edges,
+    parse_molecule, MOADDataset
+)
 from src.lightning import DDPM
 from src.linker_size_lightning import SizeClassifier
 from src.generation import generate_linkers, try_to_convert_to_sdf, get_pocket
@@ -19,7 +23,7 @@ from zipfile import ZipFile
 
 MIN_N_STEPS = 100
 MAX_N_STEPS = 500
-MAX_BATCH_SIZE = 5
+MAX_BATCH_SIZE = 20
 
 
 MODELS_METADATA = {
@@ -454,7 +458,7 @@ def generate_with_pocket(in_fragments, in_protein, n_steps, n_atoms, num_samples
     ddpm.val_dataset = dataset
 
     batch_size = min(num_samples, MAX_BATCH_SIZE)
-    dataloader = get_dataloader(dataset, batch_size=batch_size, collate_fn=collate_with_fragment_edges)
+    dataloader = get_dataloader(dataset, batch_size=batch_size, collate_fn=collate_with_fragment_without_pocket_edges)
     print('Created dataloader')
 
     ddpm.edm.T = n_steps
@@ -532,7 +536,7 @@ with demo:
                     label="Number of Denoising Steps", step=10
                 )
                 n_atoms = gr.Slider(
-                    minimum=0, maximum=20,
+                    minimum=0, maximum=20, value=5,
                     label="Linker Size: DiffLinker will predict it if set to 0",
                     step=1
                 )

src/datasets.py

@@ -317,6 +317,53 @@ def collate_with_fragment_edges(batch):
     return out
 
 
+def collate_with_fragment_without_pocket_edges(batch):
+    out = {}
+
+    # Filter out big molecules
+    # batch = [data for data in batch if data['num_atoms'] <= 50]
+
+    for i, data in enumerate(batch):
+        for key, value in data.items():
+            out.setdefault(key, []).append(value)
+
+    for key, value in out.items():
+        if key in const.DATA_LIST_ATTRS:
+            continue
+        if key in const.DATA_ATTRS_TO_PAD:
+            out[key] = torch.nn.utils.rnn.pad_sequence(value, batch_first=True, padding_value=0)
+            continue
+        raise Exception(f'Unknown batch key: {key}')
+
+    frag_mask = out['fragment_only_mask']
+    edge_mask = frag_mask[:, None, :] * frag_mask[:, :, None]
+    diag_mask = ~torch.eye(edge_mask.size(1), dtype=const.TORCH_INT, device=frag_mask.device).unsqueeze(0)
+    edge_mask *= diag_mask
+
+    batch_size, n_nodes = frag_mask.size()
+    out['edge_mask'] = edge_mask.view(batch_size * n_nodes * n_nodes, 1)
+
+    # Building edges and covalent bond values
+    rows, cols, bonds = [], [], []
+    for batch_idx in range(batch_size):
+        for i in range(n_nodes):
+            for j in range(n_nodes):
+                rows.append(i + batch_idx * n_nodes)
+                cols.append(j + batch_idx * n_nodes)
+
+    edges = [torch.LongTensor(rows).to(frag_mask.device), torch.LongTensor(cols).to(frag_mask.device)]
+    out['edges'] = edges
+
+    atom_mask = (out['fragment_mask'].bool() | out['linker_mask'].bool()).to(const.TORCH_INT)
+    out['atom_mask'] = atom_mask[:, :, None]
+
+    for key in const.DATA_ATTRS_TO_ADD_LAST_DIM:
+        if key in out.keys():
+            out[key] = out[key][:, :, None]
+
+    return out
+
+
 def get_dataloader(dataset, batch_size, collate_fn=collate, shuffle=False):
     return DataLoader(dataset, batch_size, collate_fn=collate_fn, shuffle=shuffle)
 

src/generation.py

@@ -37,7 +37,7 @@ def generate_linkers(ddpm, data, sample_fn, name, with_pocket=False, offset_idx=
     if with_pocket:
         node_mask[torch.where(data['pocket_mask'])] = 0
 
-    batch_size = len(data)
+    batch_size = len(data['positions'])
     names = [f'output_{offset_idx + i + 1}_{name}' for i in range(batch_size)]
     save_xyz_file('results', h, x, node_mask, names=names, is_geom=True, suffix='')
     print('Saved XYZ files')