Update app.py

app.py

@@ -40,32 +40,84 @@ class CFG():
     seed = 42
 
 if st.button('predict'):
-    
-    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')
+    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')
+        
+        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) 
         
-    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)
+        
+        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)
+        
+        
+        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=min_length, max_length=min_length+50, 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'])
     
-    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']
+                
+            @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',
+            )
+        
+        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=min_length, max_length=min_length+50, num_beams=CFG.num_beams, num_return_sequences=CFG.num_return_sequences, return_dict_in_generate=True, output_scores=True)
@@ -83,8 +135,7 @@ if st.button('predict'):
                     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])
@@ -92,74 +143,24 @@ if st.button('predict'):
                     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',
-        )
-    
-    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=min_length, max_length=min_length+50, 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])
+            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.append(None)
-            
+                output_df = pd.DataFrame(np.array([input_compound]+output).reshape(1, -1), columns=['input', '0th', 'valid compound'])
+            st.table(output_df)
     
-        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
-        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',
-        )
+            @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',
+            )