Update app.py

app.py

@@ -1,192 +1,179 @@
+import gc
 import os
-import random
-import numpy as np
+import sys
 import warnings
+
 import pandas as pd
-import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from torch.utils.data import Dataset, DataLoader
-import gc
 import streamlit as st
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+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 utils import seed_everything
 
 warnings.filterwarnings("ignore")
 
- 
-st.title('ReactionT5 task forward')
-st.markdown('''
+
+st.title("ReactionT5 task forward")
+st.markdown("""
 ##### At this space, you can predict the products of reactions from their inputs.
 ##### The code expects input_data as a string or CSV file that contains an "input" column. 
 ##### The format of the string or contents of the column should be "REACTANT:{reactants}REAGENT:{reagents}".
 ##### If there is no reagent, fill the blank with a space. For multiple compounds, concatenate them with ".".
 ##### The output contains SMILES of predicted products and the sum of log-likelihood for each prediction, ordered by their log-likelihood (0th is the most probable product).
-''')
-
-display_text = 'input the reaction smiles (e.g. REACTANT:COC(=O)C1=CCCN(C)C1.O.[Al+3].[H-].[Li+].[Na+].[OH-]REAGENT:C1CCOC1)'
+""")
 
 st.download_button(
-                label="Download demo_input.csv",
-                data=pd.read_csv('demo_input.csv').to_csv(index=False),
-                file_name='demo_input.csv',
-                mime='text/csv',
-            )
-
-class CFG():
-    num_beams = st.number_input(label='num beams', min_value=1, max_value=10, value=5, step=1)
-    num_return_sequences = num_beams
-    uploaded_file = st.file_uploader("Choose a CSV file")
-    input_data = st.text_area(display_text)
-    model_name_or_path = 'sagawa/ReactionT5v2-forward'
-    input_column = 'input'
-    input_max_length = 400
-    model = 't5'
-    seed = 42
-    batch_size=1
-
-def seed_everything(seed=42):
-    random.seed(seed)
-    os.environ['PYTHONHASHSEED'] = str(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
+    label="Download demo_reaction_data.csv",
+    data=pd.read_csv("data/demo_reaction_data.csv").to_csv(index=False),
+    file_name="demo_reaction_data.csv",
+    mime="text/csv",
+)
 
 
-
-def prepare_input(cfg, text):
-    inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        max_length=cfg.input_max_length,
-        padding="max_length",
-        truncation=True,
+class CFG:
+    num_beams = st.number_input(
+        label="num beams", min_value=1, max_value=10, value=5, step=1
     )
-    dic = {"input_ids": [], "attention_mask": []}
-    for k, v in inputs.items():
-        dic[k].append(torch.tensor(v[0], dtype=torch.long))
-    return dic
-
-
-class ProductDataset(Dataset):
-    def __init__(self, cfg, df):
-        self.cfg = cfg
-        self.inputs = df[cfg.input_column].values
-
-    def __len__(self):
-        return len(self.inputs)
+    num_return_sequences = num_beams
+    input_data = st.file_uploader("Choose a CSV file")
+    model_name_or_path = "sagawa/ReactionT5v2-forward"
+    input_column = "input"
+    input_max_length = 400
+    output_max_length = 300
+    model = "t5"
+    seed = 42
+    batch_size = 1
+
+
+if st.button("predict"):
+    with st.spinner(
+        "Now processing. If num beams=5, this process takes about 15 seconds per reaction."
+    ):
+        # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # seed_everything(seed=CFG.seed)
+
+        # tokenizer = AutoTokenizer.from_pretrained(CFG.model_name_or_path, return_tensors="pt")
+        # model = AutoModelForSeq2SeqLM.from_pretrained(CFG.model_name_or_path).to(device)
+        # model.eval()
+
+        # if CFG.uploaded_file is None:
+        #     input_compound = CFG.input_data
+        #     output = predict_single_input(input_compound)
+        #     sequences, scores = decode_output(output)
+        #     output_df = save_single_prediction(input_compound, sequences, scores)
+        # else:
+        #     input_data = pd.read_csv(CFG.uploaded_file)
+        #     dataset = ProductDataset(CFG, input_data)
+        #     dataloader = DataLoader(
+        #         dataset,
+        #         batch_size=CFG.batch_size,
+        #         shuffle=False,
+        #         num_workers=4,
+        #         pin_memory=True,
+        #         drop_last=False,
+        #     )
+
+        #     all_sequences, all_scores = [], []
+        #     for inputs in dataloader:
+        #         inputs = {k: v[0].to(device) for k, v in inputs.items()}
+        #         with torch.no_grad():
+        #             output = model.generate(
+        #                 **inputs,
+        #                 num_beams=CFG.num_beams,
+        #                 num_return_sequences=CFG.num_return_sequences,
+        #                 return_dict_in_generate=True,
+        #                 output_scores=True,
+        #             )
+        #         sequences, scores = decode_output(output)
+        #         all_sequences.extend(sequences)
+        #         if scores:
+        #             all_scores.extend(scores)
+        #         del output
+        #         torch.cuda.empty_cache()
+        #         gc.collect()
+
+        #     output_df = save_multiple_predictions(input_data, all_sequences, all_scores)
+
+        CFG.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        if not os.path.exists(CFG.output_dir):
+            os.makedirs(CFG.output_dir)
 
-    def __getitem__(self, idx):
-        return prepare_input(self.cfg, self.inputs[idx])
+        seed_everything(seed=CFG.seed)
 
+        CFG.tokenizer = AutoTokenizer.from_pretrained(
+            os.path.abspath(CFG.model_name_or_path)
+            if os.path.exists(CFG.model_name_or_path)
+            else CFG.model_name_or_path,
+            return_tensors="pt",
+        )
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+            os.path.abspath(CFG.model_name_or_path)
+            if os.path.exists(CFG.model_name_or_path)
+            else CFG.model_name_or_path
+        ).to(CFG.device)
+        model.eval()
 
-def predict_single_input(input_compound):
-    inp = tokenizer(input_compound, return_tensors="pt").to(device)
-    with torch.no_grad():
-        output = model.generate(
-            **inp,
-            num_beams=CFG.num_beams,
-            num_return_sequences=CFG.num_return_sequences,
-            return_dict_in_generate=True,
-            output_scores=True,
+        input_data = pd.read_csv(CFG.input_data)
+        input_data = preprocess_df(input_data, drop_duplicates=False)
+        dataset = ReactionT5Dataset(CFG, input_data)
+        dataloader = DataLoader(
+            dataset,
+            batch_size=CFG.batch_size,
+            shuffle=False,
+            num_workers=4,
+            pin_memory=True,
+            drop_last=False,
         )
-    return output
-
-
-def decode_output(output):
-    sequences = [
-        tokenizer.decode(seq, skip_special_tokens=True).replace(" ", "").rstrip(".")
-        for seq in output["sequences"]
-    ]
-    if CFG.num_beams > 1:
-        scores = output["sequences_scores"].tolist()
-        return sequences, scores
-    return sequences, None
-
-
-def save_single_prediction(input_compound, output, scores):
-    output_data = [input_compound] + output + (scores if scores else [])
-    columns = (
-        ["input"]
-        + [f"{i}th" for i in range(CFG.num_beams)]
-        + ([f"{i}th score" for i in range(CFG.num_beams)] if scores else [])
-    )
-    output_df = pd.DataFrame([output_data], columns=columns)
-    return output_df
-
-
-def save_multiple_predictions(input_data, sequences, scores):
-    output_list = [
-        [input_data.loc[i // CFG.num_return_sequences, CFG.input_column]]
-        + sequences[i : i + CFG.num_return_sequences]
-        + scores[i : i + CFG.num_return_sequences]
-        for i in range(0, len(sequences), CFG.num_return_sequences)
-    ]
-    columns = (
-        ["input"]
-        + [f"{i}th" for i in range(CFG.num_return_sequences)]
-        + ([f"{i}th score" for i in range(CFG.num_return_sequences)] if scores else [])
-    )
-    output_df = pd.DataFrame(output_list, columns=columns)
-    return output_df
 
+        all_sequences, all_scores = [], []
+        for inputs in tqdm(dataloader, total=len(dataloader)):
+            inputs = {k: v.to(CFG.device) for k, v in inputs.items()}
+            with torch.no_grad():
+                output = model.generate(
+                    **inputs,
+                    min_length=CFG.output_min_length,
+                    max_length=CFG.output_max_length,
+                    num_beams=CFG.num_beams,
+                    num_return_sequences=CFG.num_return_sequences,
+                    return_dict_in_generate=True,
+                    output_scores=True,
+                )
+            sequences, scores = decode_output(output, CFG)
+            all_sequences.extend(sequences)
+            if scores:
+                all_scores.extend(scores)
+            del output
+            torch.cuda.empty_cache()
+            gc.collect()
+
+        output_df = save_multiple_predictions(
+            input_data, all_sequences, all_scores, CFG
+        )
 
-if st.button('predict'):
-    with st.spinner('Now processing. If num beams=5, this process takes about 15 seconds per reaction.'):
+        # output_df.to_csv(os.path.join(CFG.output_dir, "output.csv"), index=False)
 
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        seed_everything(seed=CFG.seed)
-        
-        tokenizer = AutoTokenizer.from_pretrained(CFG.model_name_or_path, return_tensors="pt")
-        model = AutoModelForSeq2SeqLM.from_pretrained(CFG.model_name_or_path).to(device)
-        model.eval()
-        
-        if CFG.uploaded_file is None:
-            input_compound = CFG.input_data
-            output = predict_single_input(input_compound)
-            sequences, scores = decode_output(output)
-            output_df = save_single_prediction(input_compound, sequences, scores)
-        else:
-            input_data = pd.read_csv(CFG.uploaded_file)
-            dataset = ProductDataset(CFG, input_data)
-            dataloader = DataLoader(
-                dataset,
-                batch_size=CFG.batch_size,
-                shuffle=False,
-                num_workers=4,
-                pin_memory=True,
-                drop_last=False,
-            )
-        
-            all_sequences, all_scores = [], []
-            for inputs in dataloader:
-                inputs = {k: v[0].to(device) for k, v in inputs.items()}
-                with torch.no_grad():
-                    output = model.generate(
-                        **inputs,
-                        num_beams=CFG.num_beams,
-                        num_return_sequences=CFG.num_return_sequences,
-                        return_dict_in_generate=True,
-                        output_scores=True,
-                    )
-                sequences, scores = decode_output(output)
-                all_sequences.extend(sequences)
-                if scores:
-                    all_scores.extend(scores)
-                del output
-                torch.cuda.empty_cache()
-                gc.collect()
-        
-            output_df = save_multiple_predictions(input_data, all_sequences, all_scores)
-        
         @st.cache
         def convert_df(df):
             return df.to_csv(index=False)
-        
+
         csv = convert_df(output_df)
-        
+
         st.download_button(
             label="Download data as CSV",
             data=csv,
-            file_name='output.csv',
-            mime='text/csv',
+            file_name="output.csv",
+            mime="text/csv",
         )