Update app.py

app.py

@@ -7,19 +7,15 @@ from types import SimpleNamespace
 import pandas as pd
 import streamlit as st
 import torch
-from torch.utils.data import DataLoader
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 # Local imports
-sys.path.append(
-    os.path.abspath(os.path.join(os.path.dirname(__file__), "task_forward"))
-)
 from generation_utils import (
     ReactionT5Dataset,
     decode_output,
     save_multiple_predictions,
 )
-from train import preprocess_df
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from utils import seed_everything
 
 warnings.filterwarnings("ignore")
@@ -101,6 +97,10 @@ 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
     elif task == "retrosynthesis prediction":
         model_options = [
             "sagawa/ReactionT5v2-retrosynthesis",
@@ -109,6 +109,12 @@ 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
     else:  # yield prediction
         model_options = ["sagawa/ReactionT5v2-yield"]  # default as requested
         model_help = "Default model for yield prediction."
@@ -143,8 +149,8 @@ with st.sidebar:
         input_max_length = st.number_input(
             "Input max length",
             min_value=8,
-            max_value=input_max_length_default,
-            value=400,
+            max_value=1024,
+            value=input_max_length_default,
             step=8,
         )
         if task != "yield prediction":