alpata/ReactionPrediction

1import streamlit as st2from transformers import AutoTokenizer, AutoModelForSeq2SeqLM3from rdkit import Chem4from rdkit.Chem import Draw5from streamlit_ketcher import st_ketcher6import torch7 8# --- Page Configuration ---9st.set_page_config(10    page_title="Chemical Reaction Predictor",11    page_icon="🧪",12    layout="wide",13    initial_sidebar_state="expanded"14)15 16# --- Model Loading ---17# Use st.cache_resource to load the model only once18@st.cache_resource19def load_model():20    """21    Loads the T5 model and tokenizer from Hugging Face.22    Uses AutoModel for better compatibility.23    """24    model_name = "sagawa/ReactionT5v2-forward-USPTO_MIT"25    try:26        # Use Auto* classes for robustness27        tokenizer = AutoTokenizer.from_pretrained(model_name)28        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)29        return model, tokenizer30    except Exception as e:31        # Provide more detailed error information32        st.error("An error occurred while loading the model.")33        st.error(f"Error Type: {type(e).__name__}")34        st.error(f"Error Details: {e}")35        # Add a hint about potential memory issues on Hugging Face Spaces36        st.info("Hint: Free tiers on Hugging Face Spaces have limited memory (RAM). "37                "If the app fails to load the model, it might be due to an Out-of-Memory error. "38                "Consider upgrading your Space for more resources.")39        return None, None40 41# --- Core Functions ---42def predict_product(reactants, reagents, model, tokenizer, num_predictions):43    """Predicts the reaction product using the T5 model."""44    # Format the input string as required by the model45    # Handle the case where reagents might be empty46    if reagents and reagents.strip():47        input_text = f"reactants>{reactants}.reagents>{reagents}>products>"48    else:49        input_text = f"reactants>{reactants}>products>"50 51    input_ids = tokenizer.encode(input_text, return_tensors='pt')52 53    # Generate predictions using beam search54    outputs = model.generate(55        input_ids,56        max_length=512,57        num_beams=num_predictions * 2,  # Generate more beams for better diversity58        num_return_sequences=num_predictions,59        early_stopping=True,60    )61 62    # Decode predictions63    predictions = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]64    return predictions65 66def display_molecule(smiles_string, legend):67    """Generates and displays a molecule image from a SMILES string."""68    if not smiles_string:69        st.warning("Received an empty SMILES string.")70        return71    mol = Chem.MolFromSmiles(smiles_string)72    if mol:73        try:74            img = Draw.MolToImage(mol, size=(300, 300), legend=legend)75            st.image(img, use_column_width='auto')76        except Exception as e:77            st.warning(f"Could not generate image for SMILES: {smiles_string}. Error: {e}")78    else:79        st.warning(f"Invalid SMILES string provided: {smiles_string}")80 81# --- Initialize Session State ---82# This ensures that the state is preserved across reruns83if 'reactants' not in st.session_state:84    st.session_state.reactants = "CCO.O=C(O)C"  # Start with a default example85if 'reagents' not in st.session_state:86    st.session_state.reagents = ""87 88# --- Sidebar UI ---89with st.sidebar:90    st.title("🧪 Reaction Predictor")91    st.markdown("---")92    st.header("Controls and Information")93 94    # Example Reactions95    example_reactions = {96        "Esterification": ("CCO.O=C(O)C", ""),97        "Amide Formation": ("CCN.O=C(Cl)C", ""),98        "Suzuki Coupling": ("[B-](C1=CC=CC=C1)(F)(F)F.[K+].CC1=CC=C(Br)C=C1", "c1ccc(B(O)O)cc1"),99        "Clear Inputs": ("", "")100    }101 102    def load_example():103        # Callback to load selected example into session state104        example_key = st.session_state.example_select105        reactants, reagents = example_reactions[example_key]106        st.session_state.reactants = reactants107        st.session_state.reagents = reagents108 109    st.selectbox(110        "Load an Example Reaction",111        options=list(example_reactions.keys()),112        key="example_select",113        on_change=load_example114    )115 116    st.markdown("---")117    st.subheader("Prediction Parameters")118    num_predictions = st.slider("Number of Predictions to Generate", 1, 5, 1, help="How many potential products should the model suggest?")119    st.markdown("---")120 121    st.subheader("About")122    st.info(123        "This app uses the sagawa/ReactionT5v2-forward-USPTO_MIT model to predict chemical reaction products."124    )125    st.markdown("[View Model on Hugging Face](https://huggingface.co/sagawa/ReactionT5v2-forward-USPTO_MIT)")126 127# --- Main Application UI ---128st.title("Chemical Reaction Predictor")129st.markdown("A tool to predict chemical reactions using a state-of-the-art Transformer model.")130 131# --- Model Loading and Main Logic ---132with st.spinner("Loading the prediction model... This may take a moment on first startup."):133    model, tokenizer = load_model()134 135# Only proceed if the model loaded successfully136if model and tokenizer:137    st.success("Model loaded successfully!")138 139    # Input Section140    st.header("1. Provide Reactants and Reagents")141    input_tab1, input_tab2 = st.tabs(["✍️ Chemical Drawing Tool", "⌨️ SMILES Text Input"])142 143    with input_tab1:144        col1, col2 = st.columns(2)145        with col1:146            st.subheader("Reactants")147            # This component's value is now directly tied to the session state148            reactant_smiles_drawing = st_ketcher(st.session_state.reactants, key="ketcher_reactants")149            if reactant_smiles_drawing != st.session_state.reactants:150                st.session_state.reactants = reactant_smiles_drawing151                st.rerun() # Use the modern rerun command152 153        with col2:154            st.subheader("Reagents (Optional)")155            reagent_smiles_drawing = st_ketcher(st.session_state.reagents, key="ketcher_reagents")156            if reagent_smiles_drawing != st.session_state.reagents:157                st.session_state.reagents = reagent_smiles_drawing158                st.rerun()159 160    with input_tab2:161        st.subheader("Enter SMILES Strings")162        # Text inputs now also directly update the session state on change163        st.text_input("Reactants SMILES", key="reactant_text", value=st.session_state.reactants, on_change=lambda: setattr(st.session_state, 'reactants', st.session_state.reactant_text))164        st.text_input("Reagents SMILES", key="reagent_text", value=st.session_state.reagents, on_change=lambda: setattr(st.session_state, 'reagents', st.session_state.reagent_text))165 166    # Display the current state clearly167    st.info(f"**Current Reactants:** `{st.session_state.reactants}`")168    st.info(f"**Current Reagents:** `{st.session_state.reagents or 'None'}`")169 170    # Prediction Button171    st.header("2. Generate Prediction")172    if st.button("Predict Product", type="primary", use_container_width=True):173        if not st.session_state.reactants or not st.session_state.reactants.strip():174            st.error("Error: Reactants field cannot be empty. Please provide a molecule.")175        else:176            with st.spinner("Running prediction..."):177                predictions = predict_product(178                    st.session_state.reactants,179                    st.session_state.reagents,180                    model,181                    tokenizer,182                    num_predictions183                )184                st.header("3. Predicted Products")185                if not predictions:186                    st.warning("The model did not return any predictions.")187                else:188                    for i, product_smiles in enumerate(predictions):189                        st.subheader(f"Top Prediction #{i + 1}")190                        st.code(product_smiles, language="smiles")191                        display_molecule(product_smiles, f"Predicted Product #{i + 1}")192 193elif not model or not tokenizer:194    st.error("Application could not start because the model failed to load. Please check the error messages above.")