app.py

import os
import gc
import random
import itertools
import warnings
import logging
warnings.filterwarnings('ignore')
logging.disable(logging.WARNING)
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
import tokenizers
import transformers
from transformers import AutoTokenizer, AutoConfig, AutoModel, T5EncoderModel, get_linear_schedule_with_warmup
import datasets
from datasets import load_dataset, load_metric
import sentencepiece
import argparse
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
import torch.nn as nn
import pickle
import time
from sklearn.preprocessing import MinMaxScaler
from datasets.utils.logging import disable_progress_bar
from sklearn.metrics import mean_squared_error, r2_score
disable_progress_bar()
import streamlit as st

st.title('predictyield-t5')
st.markdown('##### At this space, you can predict the yields of reactions from their inputs.')
st.markdown('##### The code expects input_data as a string or CSV file that contains an "input" column. The format of the string or contents of the column are like "REACTANT:{reactants of the reaction}REAGENT:{reagents, catalysts, or solvents of the reaction}PRODUCT:{products of the reaction}".')
st.markdown('##### If there are no reagents or catalysts, fill the blank with a space. And if there are multiple reactants, concatenate them with "."')
display_text = 'input the reaction smiles (e.g. REACTANT:CC(C)n1ncnc1-c1cn2c(n1)-c1cnc(O)cc1OCC2.CCN(C(C)C)C(C)C.Cl.NC(=O)[C@@H]1C[C@H](F)CN1REAGENT: PRODUCT:O=C(NNC(=O)C(F)(F)F)C(F)(F)F)'


class CFG():
    uploaded_file = st.file_uploader("Choose a CSV file")
    data = st.text_area(display_text)
    pretrained_model_name_or_path = 'sagawa/ZINC-t5'
    model = 't5'
    model_name_or_path = './'
    max_len = 512
    batch_size = 5
    fc_dropout = 0.1
    seed = 42
    num_workers=1

if st.button('predict'):
    
    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)      
    
    CFG.tokenizer = AutoTokenizer.from_pretrained(CFG.model_name_or_path, return_tensors='pt')
    
    def prepare_input(cfg, text):
        inputs = cfg.tokenizer(text, add_special_tokens=True, max_length=CFG.max_len, padding='max_length', return_offsets_mapping=False, truncation=True, return_attention_mask=True)
        for k, v in inputs.items():
            inputs[k] = torch.tensor(v, dtype=torch.long)
        
        return inputs
    
    class TestDataset(Dataset):
        def __init__(self, cfg, df):
            self.cfg = cfg
            self.inputs = df['input'].values
            
        def __len__(self):
            return len(self.inputs)
        
        def __getitem__(self, item):
            inputs = prepare_input(self.cfg, self.inputs[item])
            
            return inputs
        
    
    class RegressionModel(nn.Module):
        def __init__(self, cfg, config_path=None, pretrained=False):
            super().__init__()
            self.cfg = cfg
            if config_path is None:
                self.config = AutoConfig.from_pretrained(cfg.pretrained_model_name_or_path, output_hidden_states=True)
            else:
                self.config = torch.load(config_path)
            if pretrained:
                if 't5' in cfg.model:
                    self.model = T5EncoderModel.from_pretrained(CFG.pretrained_model_name_or_path)
                else:
                    self.model = AutoModel.from_pretrained(CFG.pretrained_model_name_or_path)
            else:
                if 't5' in cfg.model:
                    self.model = T5EncoderModel.from_pretrained('sagawa/ZINC-t5')
                else:
                    self.model = AutoModel.from_config(self.config)
            self.model.resize_token_embeddings(len(cfg.tokenizer))
            self.fc_dropout1 = nn.Dropout(cfg.fc_dropout)
            self.fc1 = nn.Linear(self.config.hidden_size, self.config.hidden_size)
            self.fc_dropout2 = nn.Dropout(cfg.fc_dropout)
            self.fc2 = nn.Linear(self.config.hidden_size, 1)
            
        def forward(self, inputs):
            outputs = self.model(**inputs)
            last_hidden_states = outputs[0]
            output = self.fc1(self.fc_dropout1(last_hidden_states)[:, 0, :].view(-1, self.config.hidden_size))
            output = self.fc2(self.fc_dropout2(output))
            return output
        
    
        
    def inference_fn(test_loader, model, device):
        preds = []
        model.eval()
        model.to(device)
        tk0 = tqdm(test_loader, total=len(test_loader))
        for inputs in tk0:
            for k, v in inputs.items():
                inputs[k] = v.to(device)
            with torch.no_grad():
                y_preds = model(inputs)
            preds.append(y_preds.to('cpu').numpy())
        predictions = np.concatenate(preds)
        return predictions
    
    model = RegressionModel(CFG, config_path=CFG.model_name_or_path + '/config.pth', pretrained=False)
    state = torch.load(CFG.model_name_or_path + '/ZINC-t5_best.pth', map_location=torch.device('cpu'))
    model.load_state_dict(state)
    
    
    if CFG.uploaded_file is not None:
        test_ds = pd.read_csv(CFG.uploaded_file)
            
        test_dataset = TestDataset(CFG, test_ds)
        test_loader = DataLoader(test_dataset,
                                 batch_size=CFG.batch_size,
                                 shuffle=False,
                                 num_workers=CFG.num_workers, pin_memory=True, drop_last=False)
    
    
        prediction = inference_fn(test_loader, model, device)
    
        test_ds['prediction'] = prediction*100
        test_ds['prediction'] = test_ds['prediction'].clip(0, 100)
        csv = test_ds.to_csv(index=False)
        st.download_button(
            label="Download data as CSV",
            data=csv,
            file_name='output.csv',
            mime='text/csv'
        )
    
    else:
        test_ds = pd.DataFrame.from_dict({'input': CFG.data}, orient='index').T
        test_dataset = TestDataset(CFG, test_ds)
        test_loader = DataLoader(test_dataset,
                                 batch_size=1,
                                 shuffle=False,
                                 num_workers=CFG.num_workers, pin_memory=True, drop_last=False)
        
        
        prediction = inference_fn(test_loader, model, device)
        prediction = max(min(prediction[0][0]*100, 100), 0)
        st.text('yiled: '+ str(prediction))