output files

app.py

@@ -105,7 +105,11 @@ def generate(input_file):
         path = input_file.name
         molecule = read_molecule(path)
         name = '.'.join(path.split('/')[-1].split('.')[:-1])
-        out_sdf = f'results/{name}_generated.sdf'
+        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'
+
         print(f'Input path={path}, name={name}')
     except Exception as e:
         return f'Could not read the molecule: {e}'
@@ -113,6 +117,10 @@ def generate(input_file):
     if molecule.GetNumAtoms() > 50:
         return f'Too large molecule: upper limit is 50 heavy atoms'
 
+    with Chem.SDWriter(inp_sdf) as w:
+        w.write(molecule)
+    Chem.MolToXYZFile(molecule, inp_xyz)
+
     positions, one_hot, charges = parse_molecule(molecule, is_geom=True)
     anchors = np.zeros_like(charges)
     fragment_mask = np.ones_like(charges)
@@ -138,9 +146,9 @@ 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='generated')
+        save_xyz_file('results', h, x, node_mask, names=[name], is_geom=True, suffix='output')
         print('Saved XYZ file')
-        subprocess.run(f'obabel results/{name}_generated.xyz -O {out_sdf}', shell=True)
+        subprocess.run(f'obabel {out_xyz} -O {out_sdf}', shell=True)
         print('Converted to SDF')
         break
 
@@ -148,7 +156,7 @@ def generate(input_file):
     html = HTML_TEMPLATE.format(molecule=generated_molecule, fmt='sdf')
     return [
         IFRAME_TEMPLATE.format(html=html),
-        out_sdf,
+        [inp_sdf, inp_xyz, out_sdf, out_xyz],
     ]
 
 
@@ -165,8 +173,9 @@ with demo:
     button = gr.Button('Generate Linker!')
     
     gr.Markdown('')
+    gr.Markdown('## Output')
     visualization = gr.HTML()
-    output_files = gr.File(file_count="single", label="Output files")
+    output_files = gr.File(file_count="multiple", label="Output Files")
     
     button.click(
         fn=generate,