Update app.py

app.py

@@ -41,29 +41,48 @@ if st.button('Generate Molecules'):
     st.info("Generating molecules... Please wait.")
     # Generate molecules
     core_smiles = "C1C(=O)N(C)C(=O)C1"  # Beta-lactam core structure
+    input_ids = tokenizer(core_smiles, return_tensors='pt').input_ids
     output_ids = model.generate(
-        tokenizer(core_smiles, return_tensors='pt').input_ids,
+        input_ids=input_ids,
         max_length=128,
         temperature=creativity,
         do_sample=True,
         top_k=50,
-        num_return_sequences=num_molecules
+        num_return_sequences=num_molecules,
+        num_beams=max(num_molecules, 5)  # Ensure num_beams >= num_return_sequences
     )
     generated_smiles = [tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids]
-    molecule_names = [f"Mol{str(i).zfill(2)}" for i in range(1, num_molecules + 1)]
-    generated_molecules = dict(zip(molecule_names, generated_smiles))
+    molecule_names = [f"Mol{str(i).zfill(2)}" for i in range(1, len(generated_smiles) + 1)]
+    
+    # Create DataFrame for generated molecules
+    df_molecules = pd.DataFrame({
+        'Molecule Name': molecule_names,
+        'SMILES': generated_smiles
+    })
+    
+    # Display generated SMILES for debugging
+    st.write("Generated SMILES:")
+    st.write(df_molecules)
     
     # ADMET Predictions
-    preds = admet_model.predict(smiles=list(generated_molecules.values()))
-    preds['Molecule Name'] = molecule_names
-    preds.set_index('Molecule Name', inplace=True)
+    preds = admet_model.predict(smiles=df_molecules['SMILES'].tolist())
+    
+    # Ensure 'SMILES' is a column in preds
+    if 'SMILES' not in preds.columns:
+        preds['SMILES'] = df_molecules['SMILES']
+    
+    # Merge predictions with generated molecules
+    df_results = pd.merge(df_molecules, preds, on='SMILES', how='inner')
+    
+    # Set 'Molecule Name' as index
+    df_results.set_index('Molecule Name', inplace=True)
     
     # Display Molecules
     st.subheader('Generated Molecules')
-    cols_per_row = min(5, num_molecules)
+    cols_per_row = min(5, len(df_results))
     cols = st.columns(cols_per_row)
-    for idx, mol_name in enumerate(molecule_names):
-        smiles = generated_molecules[mol_name]
+    for idx, (mol_name, row) in enumerate(df_results.iterrows()):
+        smiles = row['SMILES']
         img = generate_molecule_image(smiles, use_safe_visualization=(string_format == 'SAFE'))
         with cols[idx % cols_per_row]:
             if isinstance(img, Image.Image):
@@ -77,10 +96,11 @@ if st.button('Generate Molecules'):
             st_copy_button(string_to_display, key=f'copy_{mol_name}')
             # Display ADMET properties
             st.write("**ADMET Properties:**")
-            st.write(preds.loc[mol_name])
+            st.write(row.drop(['SMILES']))
 else:
     st.write("Click the 'Generate Molecules' button to generate beta-lactam molecules.")
 
+
 # Function Definitions
 def generate_molecule_image(input_string, use_safe_visualization=True):
     try: