Create app.py

app.py

@@ -0,0 +1,192 @@
+import os
+import random
+import numpy as np
+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
+
+warnings.filterwarnings("ignore")
+
+ 
+st.title('ReactionT5_task_forward')
+st.markdown('''
+##### At this space, you can predict the reactants of reactions from their products.
+##### 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 smiles generated by RDKit.
+##### For multiple compounds, concatenate them with ".".
+##### The output contains SMILES of predicted reactants and the sum of log-likelihood for each prediction, ordered by their log-likelihood (0th is the most probable reactant).
+''')
+
+display_text = 'input the product smiles (e.g. CCN(CC)CCNC(=S)NC1CCCc2cc(C)cnc21)'
+
+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-retrosynthesis'
+    input_column = 'input'
+    input_max_length = 100
+    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
+
+
+
+def prepare_input(cfg, text):
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        max_length=cfg.input_max_length,
+        padding="max_length",
+        truncation=True,
+    )
+    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)
+
+    def __getitem__(self, idx):
+        return prepare_input(self.cfg, self.inputs[idx])
+
+
+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,
+        )
+    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
+
+
+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)
+        
+        @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',
+        )