Update app.py

app.py

@@ -24,14 +24,14 @@ warnings.filterwarnings("ignore")
 # Page setup
 # ------------------------------
 st.set_page_config(
-    page_title="ReactionT5 — Product Prediction",
+    page_title="ReactionT5",
     page_icon=None,
     layout="wide",
 )
 
-st.title("ReactionT5 — Product Prediction")
+st.title("ReactionT5")
 st.caption(
-    "Predict reaction products from your inputs using a pretrained ReactionT5 model."
+    "Predict reaction products, reactants, or yields from your inputs using a pretrained ReactionT5 model."
 )
 
 with st.expander("How to format your CSV", expanded=False):
@@ -97,10 +97,7 @@ with st.sidebar:
         model_help = "Recommended models for product prediction."
         input_max_length_default = 400
         output_max_length_default = 300
-        sys.path.append(
-            os.path.abspath(os.path.join(os.path.dirname(__file__), "task_forward"))
-        )
-        from train import preprocess_df
+        from task_forward.train import preprocess_df
     elif task == "retrosynthesis prediction":
         model_options = [
             "sagawa/ReactionT5v2-retrosynthesis",
@@ -109,12 +106,7 @@ with st.sidebar:
         model_help = "Recommended models for retrosynthesis prediction."
         input_max_length_default = 100
         output_max_length_default = 400
-        sys.path.append(
-            os.path.abspath(
-                os.path.join(os.path.dirname(__file__), "task_retrosynthesis")
-            )
-        )
-        from train import preprocess_df
+        from task_retrosynthesis.train import preprocess_df
     else:  # yield prediction
         model_options = ["sagawa/ReactionT5v2-yield"]  # default as requested
         model_help = "Default model for yield prediction."