Update app.py

app.py

@@ -46,7 +46,7 @@ class DrugGENConfig:
     mol_data_dir='data'
     log_dir='experiments/logs'
     model_save_dir='experiments/models'
-    inference_model=""
+    # inference_model=""
     sample_dir='experiments/samples'
     result_dir="experiments/tboard_output"
     dataset_file="chembl45_train.pt"
@@ -59,9 +59,22 @@ class DrugGENConfig:
     inf_drug_raw_file="data/akt_test.smi"
     log_sample_step=1000
 
+class ProtConfig(DrugGENConfig):
+    submodel="Prot"
+    inference_model="experiments/models/Prot"
+
+class CrossLossConfig(DrugGENConfig):
+    submodel="CrossLoss"
+    inference_model="experiments/models/CrossLoss"
+
+class NoTargetConfig(DrugGENConfig):
+    submodel="NoTarget"
+    inference_model="experiments/models/NoTarget"
+
+
     
 with st.spinner('Setting up the trainer class...'):
-    trainer = Trainer(DrugGENConfig())
+    trainer = Trainer(ProtConfig())
 
 with st.spinner('Generating Molecules...'):
     trainer.inference()