Update app.py

app.py

@@ -9,9 +9,6 @@ import time
 
 is_stopped = False
 
-seed = random.randint(0,100000)
-setup_seed(seed)
-
 def temperature_sampling(logits, temperature):
     logits = logits / temperature
     probabilities = torch.softmax(logits, dim=-1)
@@ -23,7 +20,13 @@ def stop_generation():
     is_stopped = True
     return "Generation stopped."
 
-def CTXGen(X1, X2, τ, g_num, length_range, model_name):
+def CTXGen(X1, X2, τ, g_num, length_range, model_name, seed):
+    if seed =='random':
+        seed = random.randint(0,100000)
+        setup_seed(seed)
+    else:
+        setup_seed(int(seed))
+        
     global is_stopped
     is_stopped = False
     start, end = length_range
@@ -158,6 +161,8 @@ with gr.Blocks() as demo:
     gr.Markdown("✅**Number of generations**: if it is not completed within 1200 seconds, it will automatically stop.")
     gr.Markdown("✅**Length range**: expected length range of conotoxins generated")
     gr.Markdown("✅**Model**: model parameters trained at different stages of data augmentation. Please refer to the paper for details.")
+    gr.Markdown("✅**Seed**: Enter an integer as the random seed to ensure reproducible results. The default is random")
+    
     with gr.Row():
         X1 = gr.Dropdown(choices=['<α7>','<AChBP>','<α4β2>','<α3β4>','<Ca22>','<α3β2>', '<Na12>','<α9α10>','<K16>', '<α1β1γδ>', 
                                   '<K13>', '<α1BAR>', '<α1β1ε>', '<α1AAR>', '<GluN3A>', '<GluN2B>', '<α75HT3>', '<Na14>', 
@@ -168,8 +173,10 @@ with gr.Blocks() as demo:
         X2 = gr.Dropdown(choices=['<high>','<low>'], label="Potency")
         τ = gr.Slider(minimum=1, maximum=2, step=0.1, label="τ")
         g_num = gr.Dropdown(choices=[1, 10, 20, 30, 40, 50], label="Number of generations")
+    with gr.Row():
         length_range = RangeSlider(minimum=8, maximum=50, step=1, value=(12, 16), label="Length range")
         model_name = gr.Dropdown(choices=['model_final.pt','model_C1.pt','model_C2.pt','model_C3.pt','model_C4.pt','model_C5.pt','model_mlm.pt'], label="Model")
+        seed = gr.Textbox(label="Seed", value="random")
     with gr.Row():
         start_button = gr.Button("Start Generation")
         stop_button = gr.Button("Stop Generation")
@@ -178,7 +185,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         output_df = gr.DataFrame(label="Generated Conotoxins")
     
-    start_button.click(CTXGen, inputs=[X1, X2, τ, g_num, length_range,model_name], outputs=[output_file, output_df])
+    start_button.click(CTXGen, inputs=[X1, X2, τ, g_num, length_range,model_name,seed], outputs=[output_file, output_df])
     stop_button.click(stop_generation, outputs=None)
 
 demo.launch()