Update app.py

app.py

@@ -21,6 +21,7 @@ st.markdown('### If there are no catalyst or reagent, fill the blank with a spac
 display_text = 'input the reaction smiles (e.g. REACTANT:CNc1nc(SC)ncc1CO.O.O=[Cr](=O)([O-])O[Cr](=O)(=O)[O-].[Na+]CATALYST: REAGENT: SOLVENT:CC(=O)O)'
 
 class CFG():
+    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 = 't5'
@@ -48,25 +49,66 @@ if CFG.model == 't5':
 elif CFG.model == 'deberta':
     model = EncoderDecoderModel.from_pretrained(CFG.model_name_or_path).to(device)
 
-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)
-except:
-    pass
+
+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).encode('utf-8')
+    
+    output_df = convert_df(output_df)
+    
+    st.download_button(
+        label="Download data as CSV",
+        data=output_df,
+        file_name=input_data + '_result.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)
+    except:
+        pass