sagawa commited on
Commit
2548c76
1 Parent(s): 93c5012

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +57 -0
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)