Update app.py

app.py

@@ -85,10 +85,40 @@ st.divider()
 with st.sidebar:
     st.header("Configuration")
 
-    model_name_or_path = st.text_input(
+    task = st.selectbox(
+        "Task",
+        options=["product prediction", "retrosynthesis prediction", "yield prediction"],
+        index=0,
+        help="Choose the task to run.",
+    )
+
+    # Model options tied to task
+    if task == "product prediction":
+        model_options = [
+            "sagawa/ReactionT5v2-forward",
+            "sagawa/ReactionT5v2-forward-USPTO_MIT",
+        ]
+        model_help = "Recommended models for product prediction."
+        input_max_length_default = 400
+        output_max_length_default = 300
+    elif task == "retrosynthesis prediction":
+        model_options = [
+            "sagawa/ReactionT5v2-retrosynthesis",
+            "sagawa/ReactionT5v2-retrosynthesis-USPTO_50k",
+        ]
+        model_help = "Recommended models for retrosynthesis prediction."
+        input_max_length_default = 100
+        output_max_length_default = 400
+    else:  # yield prediction
+        model_options = ["sagawa/ReactionT5v2-yield"]  # default as requested
+        model_help = "Default model for yield prediction."
+        input_max_length_default = 400
+
+    model_name_or_path = st.selectbox(
         "Model",
-        value="sagawa/ReactionT5v2-forward",
-        help="Hugging Face model repo or a local path.",
+        options=model_options,
+        index=0,
+        help=model_help,
     )
 
     num_beams = st.slider(
@@ -111,11 +141,20 @@ with st.sidebar:
 
     with st.expander("Advanced generation", expanded=False):
         input_max_length = st.number_input(
-            "Input max length", min_value=8, max_value=1024, value=400, step=8
-        )
-        output_max_length = st.number_input(
-            "Output max length", min_value=8, max_value=1024, value=300, step=8
+            "Input max length",
+            min_value=8,
+            max_value=input_max_length_default,
+            value=400,
+            step=8,
         )
+        if task != "yield prediction":
+            output_max_length = st.number_input(
+                "Output max length",
+                min_value=8,
+                max_value=1024,
+                value=output_max_length_default,
+                step=8,
+            )
         output_min_length = st.number_input(
             "Output min length",
             min_value=-1,