app.py

import streamlit as st
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import json
import os
from copy import deepcopy
import warnings
warnings.filterwarnings("ignore")

# Import custom modules
from PolymerSmilesTokenization import PolymerSmilesTokenizer
from transformers import RobertaModel, RobertaConfig, RobertaTokenizer

# Page configuration
st.set_page_config(
    page_title="TransPolymer - Polymer Property Predictor",
    page_icon="🧪",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Model definition (from Downstream.py)
class DownstreamRegression(nn.Module):
    def __init__(self, pretrained_model, drop_rate=0.1):
        super(DownstreamRegression, self).__init__()
        self.PretrainedModel = deepcopy(pretrained_model)
        
        self.Regressor = nn.Sequential(
            nn.Dropout(drop_rate),
            nn.Linear(self.PretrainedModel.config.hidden_size, self.PretrainedModel.config.hidden_size),
            nn.SiLU(),
            nn.Linear(self.PretrainedModel.config.hidden_size, 1)
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.PretrainedModel(input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs.last_hidden_state[:, 0, :]
        output = self.Regressor(logits)
        return output

# Property information
PROPERTY_INFO = {
    'PE_I': {
        'name': 'Polymer Electrolyte Conductivity I',
        'units': 'S/cm',
        'description': 'Ionic conductivity of polymer electrolytes',
        'example_smiles': 'CC(C)(C)OC(=O)NC1=CC=CC=C1'
    },
    'PE_II': {
        'name': 'Polymer Electrolyte Conductivity II',
        'units': 'S/cm', 
        'description': 'Alternative polymer electrolyte conductivity dataset',
        'example_smiles': 'CC(C)OC(=O)NC1=CC=C(C=C1)O'
    },
    'OPV': {
        'name': 'Organic Photovoltaic Efficiency',
        'units': '%',
        'description': 'Power conversion efficiency in organic solar cells',
        'example_smiles': 'c1ccc(cc1)c2ccc(cc2)C3=CC=C(C=C3)c4ccccc4'
    },
    'Egc': {
        'name': 'Band Gap (Egc)',
        'units': 'eV',
        'description': 'Electronic band gap of the polymer',
        'example_smiles': 'c1ccc2c(c1)oc1ccccc12'
    },
    'Egb': {
        'name': 'Band Gap (Egb)',
        'units': 'eV',
        'description': 'Alternative band gap measurement',
        'example_smiles': 'c1ccc2c(c1)sc1ccccc12'
    },
    'Eea': {
        'name': 'Electron Affinity',
        'units': 'eV',
        'description': 'Energy released when electron is added',
        'example_smiles': 'c1ccc2c(c1)c1ccccc1n2'
    },
    'Ei': {
        'name': 'Ionization Energy',
        'units': 'eV',
        'description': 'Energy required to remove an electron',
        'example_smiles': 'c1ccc2c(c1)c1ccccc1o2'
    },
    'Xc': {
        'name': 'Crystallization Tendency',
        'units': '-',
        'description': 'Tendency of polymer to crystallize',
        'example_smiles': 'CC(C)CC(C)(C)c1ccccc1'
    },
    'EPS': {
        'name': 'Dielectric Constant',
        'units': '-',
        'description': 'Relative permittivity of the polymer',
        'example_smiles': 'c1ccc(cc1)C(c2ccccc2)c3ccccc3'
    },
    'Nc': {
        'name': 'Refractive Index',
        'units': '-',
        'description': 'Measure of light bending in the polymer',
        'example_smiles': 'c1ccc(cc1)C#Cc2ccccc2'
    }
}

@st.cache_resource
def load_model():
    """Load the TransPolymer model (cached for performance)"""
    try:
        # Set device
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Load tokenizer
        try:
            # Try to load custom tokenizer if files exist
            if os.path.exists("ckpt/pretrain.pt") and os.path.exists("ckpt/pretrain.pt"):
                tokenizer = PolymerSmilesTokenizer.from_pretrained("roberta-base")
            else:
                # Fallback to RoBERTa base
                tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
        except:
            tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
        
        # Load model configuration
        model_config = {
            "hidden_size": 768,
            "num_hidden_layers": 6,
            "num_attention_heads": 12,
            "intermediate_size": 3072,
            "vocab_size": 50265
        }
        
        # Create model
        config = RobertaConfig(**model_config)
        pretrained_model = RobertaModel(config)
        
        # Load pretrained weights if available
        if os.path.exists("ckpt/pretrain.pt/pytorch_model.bin"):
            try:
                state_dict = torch.load("ckpt/pretrain.pt/pytorch_model.bin", map_location=device)
                pretrained_model.load_state_dict(state_dict, strict=False)
                st.success("✅ Loaded pretrained weights!")
            except:
                st.warning("⚠️ Using random weights (pretrained weights not found)")
        else:
            st.warning("⚠️ Using random weights (for demo purposes)")
        
        # Create downstream model
        model = DownstreamRegression(pretrained_model)
        model.PretrainedModel.resize_token_embeddings(len(tokenizer))
        model.to(device)
        model.eval()
        
        return model, tokenizer, device
        
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None, None, None

def predict_property(smiles, property_type, model, tokenizer, device):
    """Make prediction for a SMILES string"""
    try:
        # Tokenize SMILES
        encoded = tokenizer(
            [smiles],
            padding=True,
            truncation=True,
            max_length=411,
            return_tensors="pt"
        )
        
        # Move to device
        input_ids = encoded["input_ids"].to(device)
        attention_mask = encoded["attention_mask"].to(device)
        
        # Make prediction
        with torch.no_grad():
            outputs = model(input_ids, attention_mask)
            prediction = outputs.cpu().numpy().flatten()[0]
        
        return float(prediction)
        
    except Exception as e:
        st.error(f"Error making prediction: {e}")
        return None

def validate_smiles(smiles):
    """Basic SMILES validation"""
    if not smiles or len(smiles.strip()) == 0:
        return False, "SMILES string cannot be empty"
    
    # Basic chemical structure validation
    invalid_chars = set(smiles) - set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789()[]=#@+-.")
    if invalid_chars:
        return False, f"Invalid characters in SMILES: {invalid_chars}"
    
    return True, "Valid SMILES string"

def main():
    """Main Streamlit application"""
    
    # Header
    st.title("🧪 TransPolymer: Polymer Property Predictor")
    st.markdown("""
    **TransPolymer** is a Transformer-based language model for polymer property predictions.
    Enter a polymer SMILES string and select a property to predict its value.
    
    📄 [Paper](https://www.nature.com/articles/s41524-023-01016-5) | 
    💻 [Code](https://github.com/ChangwenXu98/TransPolymer)
    """)
    
    # Load model
    with st.spinner("Loading TransPolymer model..."):
        model, tokenizer, device = load_model()
    
    if model is None:
        st.error("Failed to load model. Please check the model files.")
        return
    
    # Sidebar for property selection
    with st.sidebar:
        st.header("🎯 Property Selection")
        
        property_type = st.selectbox(
            "Select property to predict:",
            options=list(PROPERTY_INFO.keys()),
            format_func=lambda x: f"{x}: {PROPERTY_INFO[x]['name']}"
        )
        
        # Display property information
        prop_info = PROPERTY_INFO[property_type]
        st.info(f"""
        **{prop_info['name']}**
        
        {prop_info['description']}
        
        **Units:** {prop_info['units']}
        """)
        
        # Example SMILES
        st.subheader("📝 Example")
        if st.button("Use Example SMILES"):
            st.session_state.example_smiles = prop_info['example_smiles']
    
    # Main interface
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.header("📮 Input")
        
        # SMILES input
        smiles_input = st.text_area(
            "Enter polymer SMILES string:",
            value=st.session_state.get('example_smiles', ''),
            height=100,
            placeholder="e.g., CC(C)(C)OC(=O)NC1=CC=CC=C1",
            help="SMILES (Simplified Molecular Input Line Entry System) notation for the polymer structure"
        )
        
        # Validation
        if smiles_input:
            is_valid, message = validate_smiles(smiles_input)
            if is_valid:
                st.success(f"✅ {message}")
            else:
                st.error(f"❌ {message}")
        
        # Predict button
        predict_button = st.button("🔮 Predict Property", type="primary", disabled=not smiles_input)
    
    with col2:
        st.header("📊 Results")
        
        if predict_button and smiles_input:
            with st.spinner("Making prediction..."):
                prediction = predict_property(smiles_input, property_type, model, tokenizer, device)
            
            if prediction is not None:
                prop_info = PROPERTY_INFO[property_type]
                
                # Display prediction
                st.metric(
                    label=f"{prop_info['name']}",
                    value=f"{prediction:.2e}" if abs(prediction) < 0.01 else f"{prediction:.3f}",
                    delta=None
                )
                
                st.info(f"**Units:** {prop_info['units']}")
                
                # Additional info
                with st.expander("📋 Prediction Details"):
                    st.json({
                        "smiles": smiles_input,
                        "property": property_type,
                        "prediction": float(prediction),
                        "units": prop_info['units'],
                        "model": "TransPolymer",
                        "note": "This is a research model for demonstration purposes"
                    })
    
    # Additional sections
    st.markdown("---")
    
    # Batch prediction
    with st.expander("📊 Batch Predictions"):
        st.markdown("Upload a CSV file with SMILES strings for batch prediction:")
        
        uploaded_file = st.file_uploader("Choose CSV file", type="csv")
        
        if uploaded_file is not None:
            try:
                df = pd.read_csv(uploaded_file)
                st.write("Preview of uploaded data:")
                st.dataframe(df.head())
                
                if 'smiles' in df.columns:
                    if st.button("Run Batch Prediction"):
                        with st.spinner("Processing batch predictions..."):
                            predictions = []
                            progress_bar = st.progress(0)
                            
                            for i, smiles in enumerate(df['smiles']):
                                pred = predict_property(smiles, property_type, model, tokenizer, device)
                                predictions.append(pred)
                                progress_bar.progress((i + 1) / len(df))
                            
                            df[f'predicted_{property_type}'] = predictions
                            st.success("✅ Batch prediction completed!")
                            st.dataframe(df)
                            
                            # Download results
                            csv = df.to_csv(index=False)
                            st.download_button(
                                label="📥 Download Results",
                                data=csv,
                                file_name=f"transpolymer_predictions_{property_type}.csv",
                                mime="text/csv"
                            )
                else:
                    st.error("CSV file must contain a 'smiles' column")
                    
            except Exception as e:
                st.error(f"Error processing file: {e}")
    
    # About section
    with st.expander("ℹ️ About TransPolymer"):
        st.markdown("""
        **TransPolymer** is a Transformer-based language model for polymer property predictions developed by researchers at Carnegie Mellon University.
        
        **Key Features:**
        - 6-layer RoBERTa architecture with 768 hidden dimensions
        - Pretrained on ~5 million polymer SMILES sequences
        - Custom chemical-aware tokenization
        - Supports 10+ different polymer properties
        
        **Citation:**
        ```Q
        @article{xu2023transpolymer,
          title={TransPolymer: a Transformer-based language model for polymer property predictions},
          author={Xu, Changwen and Wang, Yuyang and Barati Farimani, Amir},
          journal={npj Computational Materials},
          volume={9},
          number={1},
          pages={64},
          year={2023},
          publisher={Nature Publishing Group UK London}
        }
        ```
        
        **Disclaimer:** This is a research model for demonstration purposes. Results should be validated experimentally.
        """)

if __name__ == "__main__":
    main()