Update app.py

app.py

@@ -1,24 +1,22 @@
 import os
-import gc
 import random
-import warnings
-warnings.filterwarnings('ignore')
 import numpy as np
+import warnings
 import pandas as pd
 import torch
-import tokenizers
-import transformers
-from transformers import AutoTokenizer, EncoderDecoderModel, AutoModelForSeq2SeqLM
-import sentencepiece
-from rdkit import Chem
-import rdkit
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from torch.utils.data import Dataset, DataLoader
+import gc
 import streamlit as st
+
+warnings.filterwarnings("ignore")
+
  
-st.title('predictproduct-t5')
+st.title('ReactionT5_task_forward')
 st.markdown('##### At this space, you can predict the products of reactions from their inputs.')
 st.markdown('##### 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 are like "REACTANT:{reactants of the reaction}REAGENT:{reagents, catalysts, or solvents of the reaction}".')
 st.markdown('##### If there is no reagent, fill the blank with a space. And if there are multiple compounds, concatenate them with "."')
-st.markdown('##### The output contains smiles of predicted products and sum of log-likelihood for each prediction. Predictions are ordered by their log-likelihood.(0th is the most probable product.) "valid compound" is the most probable and valid(can be recognized by RDKit) prediction.')
+st.markdown('##### The output contains smiles of predicted products and sum of log-likelihood for each prediction. Predictions are 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)'
@@ -35,120 +33,150 @@ class CFG():
     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/ZINC-t5-productpredicition'
+    model_name_or_path = 'sagawa/ReactionT5-forward-v2'
+    input_column = 'input'
+    input_max_length = 400
     model = 't5'
     seed = 42
 
+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")
         
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        
-        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
-        seed_everything(seed=CFG.seed) 
-        
-        
-        tokenizer = AutoTokenizer.from_pretrained(CFG.model_name_or_path, return_tensors='pt')
-            
-        if CFG.model == 't5':
-            model = AutoModelForSeq2SeqLM.from_pretrained(CFG.model_name_or_path).to(device)
-        elif CFG.model == 'deberta':
-            model = EncoderDecoderModel.from_pretrained(CFG.model_name_or_path).to(device)
+        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 not None:
-            input_data = pd.read_csv(CFG.uploaded_file)
-            outputs = []
-            for idx, row in input_data.iterrows():
-                input_compound = row['input']
-                # min_length = min(input_compound.find('CATALYST') - input_compound.find(':') - 10, 0)
-                inp = tokenizer(input_compound, return_tensors='pt').to(device)
-                output = model.generate(**inp, min_length=2, max_length=181, num_beams=CFG.num_beams, num_return_sequences=CFG.num_return_sequences, return_dict_in_generate=True, output_scores=True)
-                if CFG.num_beams > 1:
-                    scores = output['sequences_scores'].tolist()
-                    output = [tokenizer.decode(i, skip_special_tokens=True).replace(' ', '').rstrip('.') for i in output['sequences']]
-                    for ith, out in enumerate(output):
-                        mol = Chem.MolFromSmiles(out.rstrip('.'))
-                        if type(mol) == rdkit.Chem.rdchem.Mol:
-                            output.append(out.rstrip('.'))
-                            scores.append(scores[ith])
-                            break
-                    if type(mol) == None:
-                        output.append(None)
-                        scores.append(None)
-                    output += scores
-                    output = [input_compound] + output
-                    outputs.append(output)
-    
-                else:
-                    output = [tokenizer.decode(output['sequences'][0], skip_special_tokens=True).replace('. ', '.').rstrip('.')]
-                    mol = Chem.MolFromSmiles(output[0])
-                    if type(mol) == rdkit.Chem.rdchem.Mol:
-                        output.append(output[0])
-                    else:
-                        output.append(None)
-                    output = [input_compound] + output
-                    outputs.append(output)    
-                    
-            if CFG.num_beams > 1:
-                output_df = pd.DataFrame(outputs, columns=['input'] + [f'{i}th' for i in range(CFG.num_beams)] + ['valid compound'] + [f'{i}th score' for i in range(CFG.num_beams)] + ['valid compound score'])
-            else:
-                output_df = pd.DataFrame(outputs, columns=['input', '0th', 'valid compound'])
-    
-                
-            @st.cache
-            def convert_df(df):
-                # IMPORTANT: Cache the conversion to prevent computation on every rerun
-                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',
+        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.input_data)
+            dataset = ProductDataset(CFG, input_data)
+            dataloader = DataLoader(
+                dataset,
+                batch_size=CFG.batch_size,
+                shuffle=False,
+                num_workers=4,
+                pin_memory=True,
+                drop_last=False,
             )
         
-        else:
-            input_compound = CFG.input_data
-            # min_length = min(input_compound.find('CATALYST') - input_compound.find(':') - 10, 0)
-            inp = tokenizer(input_compound, return_tensors='pt').to(device)
-            output = model.generate(**inp, min_length=2, max_length=181, num_beams=CFG.num_beams, num_return_sequences=CFG.num_return_sequences, return_dict_in_generate=True, output_scores=True)
-            if CFG.num_beams > 1:
-                scores = output['sequences_scores'].tolist()
-                output = [tokenizer.decode(i, skip_special_tokens=True).replace(' ', '').rstrip('.') for i in output['sequences']]
-                for ith, out in enumerate(output):
-                    mol = Chem.MolFromSmiles(out.rstrip('.'))
-                    if type(mol) == rdkit.Chem.rdchem.Mol:
-                        output.append(out.rstrip('.'))
-                        scores.append(scores[ith])
-                        break
-                if type(mol) == None:
-                    output.append(None)
-                    scores.append(None)
-                output += scores
-                output = [input_compound] + output
-                
-            else:
-                output = [tokenizer.decode(output['sequences'][0], skip_special_tokens=True).replace('. ', '.').rstrip('.')]
-                mol = Chem.MolFromSmiles(output[0])
-                if type(mol) == rdkit.Chem.rdchem.Mol:
-                    output.append(output[0])
-                else:
-                    output.append(None)
-                
+            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,
+                        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)
+                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)
         
-            if CFG.num_beams > 1:
-                output_df = pd.DataFrame(np.array(output).reshape(1, -1), columns=['input'] + [f'{i}th' for i in range(CFG.num_beams)] + ['valid compound'] + [f'{i}th score' for i in range(CFG.num_beams)] + ['valid compound score'])
-            else:
-                output_df = pd.DataFrame(np.array([input_compound]+output).reshape(1, -1), columns=['input', '0th', 'valid compound'])
             st.table(output_df)
     
             @st.cache