Update app.py

app.py

@@ -321,16 +321,55 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
                 info="Choose which protein target or general model to use for molecule generation"
             )
             
-            input_mode = gr.Radio(
-                choices=["generate", "smiles"],
-                value="generate",
-                label="Input Mode",
-                info="Choose to generate new molecules or provide your own SMILES strings",
-                elem_id="input_mode"
-            )
+            # Add a separator between model selection and input mode
+            gr.Markdown("---")
+            gr.Markdown("## Input Settings")
+            
+            # Replace radio with switch using a better layout
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=1, min_width=150):
+                    gr.Markdown("### Classic Generation", elem_id="generate-mode-label")
+                
+                with gr.Column(scale=1, min_width=150):
+                    input_mode_switch = gr.Switch(
+                        value=False,
+                        label="Custom SMILES Input",
+                        elem_id="input-mode-switch"
+                    )
+                    
+                with gr.Column(scale=1, min_width=150):
+                    gr.Markdown("### Custom SMILES Input", elem_id="smiles-mode-label")
+            
+            # Add custom CSS for the switch styling
+            gr.HTML("""
+            <style>
+            #input-mode-switch {
+                margin: 0 auto;
+                display: flex;
+                justify-content: center;
+                transform: scale(1.5);
+            }
+            #generate-mode-label, #smiles-mode-label {
+                text-align: center;
+                margin-top: 10px;
+            }
+            .input-box {
+                border: 2px solid rgba(128, 128, 228, 0.3);
+                border-radius: 10px;
+                padding: 15px;
+                margin-top: 15px;
+                background-color: rgba(32, 36, 45, 0.7);
+                box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+            }
+            .input-box:hover {
+                border-color: rgba(128, 128, 228, 0.6);
+                box-shadow: 0 6px 8px rgba(0, 0, 0, 0.15);
+            }
+            </style>
+            """)
             
             # Create container for generation mode inputs
-            with gr.Group(visible=True) as generate_group:
+            with gr.Box(visible=True, elem_id="generate-box", elem_classes="input-box") as generate_group:
                 num_molecules = gr.Slider(
                     minimum=10,
                     maximum=250,
@@ -339,27 +378,37 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
                     label="Number of Molecules to Generate",
                     info="This space runs on a CPU, which may result in slower performance. Generating 200 molecules takes approximately 6 minutes. Therefore, We set a 250-molecule cap. On a GPU, the model can generate 10,000 molecules in the same amount of time. Please check our GitHub repo for running our models on GPU."
                 )
+                
+                # Seed input used in generate mode
+                seed_num_generate = gr.Textbox(
+                    label="Random Seed (Optional)",
+                    value="",
+                    info="Set a specific seed for reproducible results, or leave empty for random generation"
+                )
             
             # Create container for SMILES input mode
-            with gr.Group(visible=False) as smiles_group:
+            with gr.Box(visible=False, elem_id="smiles-box", elem_classes="input-box") as smiles_group:
                 smiles_input = gr.Textbox(
                     label="Input SMILES",
                     info="Enter up to 100 SMILES strings, one per line",
                     lines=10,
                     placeholder="CC(=O)OC1=CC=CC=C1C(=O)O\nCCO\nC1=CC=C(C=C1)C(=O)O\n...",
                 )
-            
-            # Seed input is used by both modes
-            seed_num = gr.Textbox(
-                label="Random Seed (Optional)",
-                value="",
-                info="Set a specific seed for reproducible results, or leave empty for random generation"
-            )
+                
+                # Seed input used in SMILES mode
+                seed_num_smiles = gr.Textbox(
+                    label="Random Seed (Optional)",
+                    value="",
+                    info="Set a specific seed for reproducible results, or leave empty for random generation"
+                )
 
             # Handle visibility toggling between the two input modes
-            input_mode.change(
-                fn=lambda x: [x == "generate", x == "smiles"],
-                inputs=[input_mode],
+            def toggle_visibility(checkbox_value):
+                return not checkbox_value, checkbox_value
+                
+            input_mode_switch.change(
+                fn=toggle_visibility,
+                inputs=[input_mode_switch],
                 outputs=[generate_group, smiles_group]
             )
 
@@ -368,6 +417,12 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
                 variant="primary",
                 size="lg"
             )
+            
+            # Helper function to determine which mode is active and which seed to use
+            def get_inputs(checkbox_value, num_mols, seed_gen, seed_smiles, smiles):
+                mode = "smiles" if checkbox_value else "generate"
+                seed = seed_smiles if checkbox_value else seed_gen
+                return [mode, num_mols, seed, smiles]
 
         with gr.Column(scale=2):
             basic_metrics_df = gr.Dataframe(
@@ -393,8 +448,14 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
     gr.Markdown("### Created by the HUBioDataLab | [GitHub](https://github.com/HUBioDataLab/DrugGEN) | [Paper](https://arxiv.org/abs/2302.07868)")
 
     submit_button.click(
-        function, 
-        inputs=[model_name, input_mode, num_molecules, seed_num, smiles_input], 
+        fn=lambda model, checkbox, num_mols, seed_gen, seed_smiles, smiles: function(
+            model, 
+            "smiles" if checkbox else "generate", 
+            num_mols, 
+            seed_smiles if checkbox else seed_gen, 
+            smiles
+        ),
+        inputs=[model_name, input_mode_switch, num_molecules, seed_num_generate, seed_num_smiles, smiles_input], 
         outputs=[
             image_output, 
             file_download,