multiple samples

app.py

@@ -14,6 +14,8 @@ from src.datasets import get_dataloader, collate_with_fragment_edges, parse_mole
 from src.lightning import DDPM
 from src.linker_size_lightning import SizeClassifier
 
+N_SAMPLES = 5
+
 parser = argparse.ArgumentParser()
 parser.add_argument('--ip', type=str, default=None)
 args = parser.parse_args()
@@ -103,10 +105,8 @@ def generate(input_file):
         molecule = read_molecule(path)
         molecule = Chem.RemoveAllHs(molecule)
         name = '.'.join(path.split('/')[-1].split('.')[:-1])
-        inp_sdf = f'results/{name}_input.sdf'
-        inp_xyz = f'results/{name}_input.xyz'
-        out_sdf = f'results/{name}_output.sdf'
-        out_xyz = f'results/{name}_output.xyz'
+        inp_sdf = f'results/input_{name}.sdf'
+        inp_xyz = f'results/input_{name}.xyz'
     except Exception as e:
         return f'Could not read the molecule: {e}'
 
@@ -133,8 +133,8 @@ def generate(input_file):
         'fragment_mask': torch.tensor(fragment_mask, dtype=const.TORCH_FLOAT, device=device),
         'linker_mask': torch.tensor(linker_mask, dtype=const.TORCH_FLOAT, device=device),
         'num_atoms': len(positions),
-    }]
-    dataloader = get_dataloader(dataset, batch_size=1, collate_fn=collate_with_fragment_edges)
+    }] * N_SAMPLES
+    dataloader = get_dataloader(dataset, batch_size=N_SAMPLES, collate_fn=collate_with_fragment_edges)
     print('Created dataloader')
 
     for data in dataloader:
@@ -142,12 +142,21 @@ def generate(input_file):
         print('Generated linker')
         x = chain[0][:, :, :ddpm.n_dims]
         h = chain[0][:, :, ddpm.n_dims:]
-        save_xyz_file('results', h, x, node_mask, names=[name], is_geom=True, suffix='output')
-        print('Saved XYZ file')
-        subprocess.run(f'obabel {out_xyz} -O {out_sdf}', shell=True)
-        print('Converted to SDF')
+        names = [f'output_{i+1}_{name}' for i in range(N_SAMPLES)]
+        save_xyz_file('results', h, x, node_mask, names=names, is_geom=True, suffix='')
+        print('Saved XYZ files')
         break
 
+    out_files = []
+    for i in range(N_SAMPLES):
+        out_xyz = f'results/output_{i+1}_{name}_.xyz'
+        out_sdf = f'results/output_{i+1}_{name}_.sdf'
+        subprocess.run(f'obabel {out_xyz} -O {out_sdf}', shell=True)
+        out_files.append(out_xyz)
+        out_files.append(out_sdf)
+    print('Converted to SDF')
+
+    out_sdf = f'results/output_1_{name}_.sdf'
     input_fragments_content = read_molecule_content(inp_sdf)
     generated_molecule_content = read_molecule_content(out_sdf)
     html = output.SAMPLES_RENDERING_TEMPLATE.format(
@@ -158,7 +167,7 @@ def generate(input_file):
     )
     return [
         output.IFRAME_TEMPLATE.format(html=html),
-        [inp_sdf, inp_xyz, out_sdf, out_xyz],
+        [inp_sdf, inp_xyz] + out_files,
     ]