Update app.py

app.py

@@ -94,12 +94,19 @@ with st.sidebar:
         "Select a model to make inference",
         ('Prot', 'CrossLoss', 'NoTarget'))
 
+        molecule_num = st.number_input('Number of molecules to generate', min_value=1, max_value=100_000, value=1000, step=1)
+
         submitted = st.form_submit_button("Start Computing")
 
 
 if submitted:
+
+    config = model_configs[model_name]
+
+    config.inference_sample_num = molecule_num
+    
     with st.spinner(f'Creating the trainer class instance for {model_name}...'):
-        trainer = Trainer(model_configs[model_name])
+        trainer = Trainer(config)
     with st.spinner(f'Running inference function of {model_name} (this may take a while) ...'):
         trainer.inference()
     st.success(f"Success with the inference of {model_name}")