Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,57 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import gc
|
3 |
+
import random
|
4 |
+
import warnings
|
5 |
+
warnings.filterwarnings('ignore')
|
6 |
+
import numpy as np
|
7 |
+
import pandas as pd
|
8 |
+
import torch
|
9 |
+
import tokenizers
|
10 |
+
import transformers
|
11 |
+
from transformers import AutoTokenizer, EncoderDecoderModel, AutoModelForSeq2SeqLM
|
12 |
+
import sentencepiece
|
13 |
+
from rdkit import Chem
|
14 |
+
import rdkit
|
15 |
+
import streamlit as st
|
16 |
+
|
17 |
+
class CFG():
|
18 |
+
input_data = st.text_area('enter chemical reaction (e.g. REACTANT:NCCO.O=C1COCC(=O)O1CATALYST: REAGENT: SOLVENT:c1ccncc1)')
|
19 |
+
model_name_or_path = 'sagawa/ZINC-t5'
|
20 |
+
model = 't5'
|
21 |
+
num_beams = 5
|
22 |
+
num_return_sequences = 5
|
23 |
+
seed = 42
|
24 |
+
|
25 |
+
|
26 |
+
|
27 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
28 |
+
|
29 |
+
def seed_everything(seed=42):
|
30 |
+
random.seed(seed)
|
31 |
+
os.environ['PYTHONHASHSEED'] = str(seed)
|
32 |
+
np.random.seed(seed)
|
33 |
+
torch.manual_seed(seed)
|
34 |
+
torch.cuda.manual_seed(seed)
|
35 |
+
torch.backends.cudnn.deterministic = True
|
36 |
+
seed_everything(seed=CFG.seed)
|
37 |
+
|
38 |
+
|
39 |
+
input_compound = CFG.input_data
|
40 |
+
min_length = min(input_compound.find('CATALYST') - input_compound.find(':') - 10, 0)
|
41 |
+
inp = tokenizer(input_compound, return_tensors='pt').to(device)
|
42 |
+
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)
|
43 |
+
scores = output['sequences_scores'].tolist()
|
44 |
+
output = [tokenizer.decode(i, skip_special_tokens=True).replace('. ', '.').rstrip('.') for i in output['sequences']]
|
45 |
+
for ith, out in enumerate(output):
|
46 |
+
mol = Chem.MolFromSmiles(out.rstrip('.'))
|
47 |
+
if type(mol) == rdkit.Chem.rdchem.Mol:
|
48 |
+
output.append(out.rstrip('.'))
|
49 |
+
scores.append(scores[ith])
|
50 |
+
break
|
51 |
+
if type(mol) == None:
|
52 |
+
output.append(None)
|
53 |
+
scores.append(None)
|
54 |
+
output += scores
|
55 |
+
output = [input_compound] + output
|
56 |
+
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'])
|
57 |
+
print(output_df)
|