Update app.py

app.py

@@ -62,23 +62,38 @@ if st.button('predict'):
             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 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(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'])
-        
+                    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
@@ -98,21 +113,34 @@ if st.button('predict'):
         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 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(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'])
+                output.append(output[0])
+            else:
+                output.append(None)
+            
+    
+            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