import gradio as gr
import pandas as pd
from rdkit import Chem
from rdkit.Chem import AllChem
import joblib

model = joblib.load('CHO.pkl')

def predict(smiles):
    if smiles.strip() == "":
        raise gr.Error("SMILES input error")

    mol = Chem.MolFromSmiles(smiles)
    if mol == None:
        raise gr.Error("SMILES input error")
    mol_ECFP4 = list(AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=1024).ToBitString())
    preprocess_data = pd.DataFrame([mol_ECFP4])
    result = model.predict(preprocess_data)
    postprocess_data = '{:.2e}'.format(pow(10, result[0]))
    return postprocess_data

iface = gr.Interface(fn=predict, inputs=gr.Textbox(lines=2, label="Chemical substance SMILES"), 
                     outputs=gr.Textbox(lines=1, label="Cytotoxicity of disinfection byproducts in CHO cells",info="Unit of measurement: molar concentration"),
                     examples=[["O=C(O)CBr"],["O=CC(Br)(Br)Br"],["IC(Br)Br"]],allow_flagging="never")
iface.launch()