Update app.py

app.py

@@ -29,32 +29,69 @@ class CFG():
     model = 't5'
     seed = 42
 
-
-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 st.button('predict'):
     
-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']
+    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)
+    
+    
+    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)
+            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)
+    
+        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'])
+        
+        @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)
@@ -71,59 +108,23 @@ if CFG.uploaded_file is not None:
             scores.append(None)
         output += scores
         output = [input_compound] + output
-        outputs.append(output)
-
-    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'])
+        try:
+            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'])
+            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.cache
+            def convert_df(df):
+                # IMPORTANT: Cache the conversion to prevent computation on every rerun
+                return df.to_csv(index=False)
     
-    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)
-    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
-    try:
-        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'])
-        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',
-        )
-        
-    except:
-        pass
+            csv = convert_df(output_df)
+            
+            st.download_button(
+                label="Download data as CSV",
+                data=csv,
+                file_name='output.csv',
+                mime='text/csv',
+            )
+            
+        except:
+            pass