Update gradio_app.py

gradio_app.py

@@ -86,8 +86,8 @@ class NoTargetConfig(DrugGENConfig):
 
 model_configs = {
     "Prot": ProtConfig(),
-    "CrossLoss": CrossLossConfig(),
-    "NoTarget": NoTargetConfig(),
+    "DrugGEN": CrossLossConfig(),
+    "DrugGEN-NoTarget": NoTargetConfig(),
 }
 
 
@@ -97,7 +97,6 @@ def function(model_name: str, mol_num: int, seed: int) -> tuple[PIL.Image, pd.Da
     Returns:
     image, score_df, file path
     '''
-    model_name = model_name.replace("DrugGEN-", "")
 
     config = model_configs[model_name]
     config.inference_sample_num = mol_num
@@ -160,8 +159,8 @@ with gr.Blocks() as demo:
 - **DrugGEN-NoTarget**: composed of one GAN, focuses on learning the chemical properties from the ChEMBL training dataset, no target-specific generation.
         """)
             model_name = gr.Radio(
-                choices=("DrugGEN-Prot", "DrugGEN-CrossLoss", "DrugGEN-NoTarget"),
-                value="DrugGEN-Prot",
+                choices=("DrugGEN", "DrugGEN-NoTarget"),
+                value="DrugGEN",
                 label="Select a model to make inference",
                 info=" DrugGEN-Prot and DrugGEN-CrossLoss models design molecules to target the AKT1 protein"
             )