Update app.py

app.py

@@ -1,680 +1,680 @@
-# app.py
-
-import streamlit as st
-import torch
-import torch.nn as nn
-import pickle
-import numpy as np
-import pandas as pd
-from typing import List, Dict, Tuple, Optional
-
-# RDKit for molecule handling
-from rdkit import Chem
-from rdkit.Chem import Draw, Descriptors
-from rdkit import RDLogger
-RDLogger.DisableLog('rdApp.*')
-
-# Visualization libraries
-import matplotlib.pyplot as plt
-import seaborn as sns
-
-# For generating images in Streamlit
-from PIL import Image
-
-# Suppress warnings in RDKit
-import warnings
-warnings.filterwarnings('ignore')
-
-# Set Seaborn style
-sns.set_style('whitegrid')
-
-# Additional imports for GNN
-import torch.nn.functional as F
-from torch.nn import Linear, Sequential, BatchNorm1d, ReLU
-
-from torch_geometric.data import Data
-from torch_geometric.nn import GCNConv, GINConv
-from torch_geometric.nn import global_mean_pool, global_add_pool
-
-# Function to load the VAE model
-@st.cache_resource
-def load_vae_model(device):
-    # Load the vocabulary
-    with open('vae_vocab.pkl', 'rb') as f:
-        vocab = pickle.load(f)
-    vocab_size = len(vocab)
-    
-    # Initialize the model with the same parameters
-    hidden_dim = 256  # Ensure this matches your trained model
-    latent_dim = 64   # Ensure this matches your trained model
-    
-    # Define the VAE class (same as in your training script)
-    class VAE(nn.Module):
-        def __init__(self, vocab_size: int, hidden_dim: int, latent_dim: int):
-            super(VAE, self).__init__()
-            self.vocab_size = vocab_size
-            self.hidden_dim = hidden_dim
-            self.latent_dim = latent_dim
-
-            self.encoder = nn.GRU(vocab_size, hidden_dim, batch_first=True)
-            self.fc_mu = nn.Linear(hidden_dim, latent_dim)
-            self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
-
-            self.decoder = nn.GRU(vocab_size + latent_dim, hidden_dim, batch_first=True)
-            self.fc_output = nn.Linear(hidden_dim, vocab_size)
-        
-        def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-            _, h = self.encoder(x)
-            h = h.squeeze(0)
-            return self.fc_mu(h), self.fc_logvar(h)
-
-        def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
-            std = torch.exp(0.5 * logvar)
-            eps = torch.randn_like(std)
-            return mu + eps * std
-
-        def decode(self, z: torch.Tensor, max_length: int) -> torch.Tensor:
-            batch_size = z.size(0)
-            h = torch.zeros(1, batch_size, self.hidden_dim).to(z.device)
-            x = torch.zeros(batch_size, 1, self.vocab_size).to(z.device)
-            x[:, 0, vocab['<']] = 1  # Start token
-            outputs = []
-
-            for _ in range(max_length):
-                z_input = z.unsqueeze(1)
-                decoder_input = torch.cat([x, z_input], dim=2)
-                output, h = self.decoder(decoder_input, h)
-                output = self.fc_output(output)
-                outputs.append(output)
-                x = torch.softmax(output, dim=-1)
-
-            return torch.cat(outputs, dim=1)
-
-        def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            mu, logvar = self.encode(x)
-            z = self.reparameterize(mu, logvar)
-            return self.decode(z, x.size(1)), mu, logvar
-
-    model = VAE(vocab_size, hidden_dim, latent_dim)
-    model.load_state_dict(torch.load('vae_model.pth', map_location=device))
-    model.to(device)
-    model.eval()
-    return model, vocab
-
-# Function to generate molecules using VAE
-def generate_smiles_vae(model, vocab, num_samples=10, max_length=100):
-    model.eval()
-    inv_vocab = {v: k for k, v in vocab.items()}
-    generated_smiles = []
-    device = next(model.parameters()).device
-
-    with torch.no_grad():
-        for _ in range(num_samples):
-            z = torch.randn(1, model.latent_dim).to(device)
-            x = torch.zeros(1, 1, model.vocab_size).to(device)
-            x[0, 0, vocab['<']] = 1
-            h = torch.zeros(1, 1, model.hidden_dim).to(device)
-
-            smiles = ''
-            for _ in range(max_length):
-                z_input = z.unsqueeze(1)
-                decoder_input = torch.cat([x, z_input], dim=2)
-                output, h = model.decoder(decoder_input, h)
-                output = model.fc_output(output)
-
-                probs = torch.softmax(output.squeeze(0), dim=-1)
-                next_char = torch.multinomial(probs, 1).item()
-
-                if next_char == vocab['>']:
-                    break
-
-                smiles += inv_vocab.get(next_char, '')
-                x = torch.zeros(1, 1, model.vocab_size).to(device)
-                x[0, 0, next_char] = 1
-
-            generated_smiles.append(smiles)
-
-    return generated_smiles
-
-# Function to post-process and validate SMILES strings
-def enhanced_post_process_smiles(smiles: str) -> str:
-    smiles = smiles.replace('<', '').replace('>', '')
-    allowed_chars = set('CNOPSFIBrClcnops()[]=@+-#0123456789')
-    smiles = ''.join(c for c in smiles if c in allowed_chars)
-
-    # Balance parentheses
-    open_count = smiles.count('(')
-    close_count = smiles.count(')')
-    if open_count > close_count:
-        smiles += ')' * (open_count - close_count)
-    elif close_count > open_count:
-        smiles = '(' * (close_count - open_count) + smiles
-
-    # Replace invalid double bonds
-    smiles = smiles.replace('==', '=')
-
-    # Attempt to close unclosed rings
-    for i in range(1, 10):
-        if smiles.count(str(i)) % 2 != 0:
-            smiles += str(i)
-
-    return smiles
-
-def validate_and_correct_smiles(smiles: str) -> Tuple[bool, str]:
-    mol = Chem.MolFromSmiles(smiles)
-    if mol is not None:
-        try:
-            Chem.SanitizeMol(mol)
-            return True, Chem.MolToSmiles(mol, isomericSmiles=True)
-        except:
-            pass
-    return False, smiles
-
-# Function to analyze molecules
-def analyze_molecules(smiles_list: List[str], training_smiles_set: set) -> Dict:
-    results = {
-        'total': len(smiles_list),
-        'valid': 0,
-        'invalid': 0,
-        'unique': 0,
-        'corrected': 0,
-        'novel': 0,
-        'valid_properties': [],
-        'novel_properties': [],
-        'invalid_smiles': []
-    }
-
-    unique_smiles = set()
-    novel_smiles = set()
-
-    for smiles in smiles_list:
-        processed_smiles = enhanced_post_process_smiles(smiles)
-        is_valid, corrected_smiles = validate_and_correct_smiles(processed_smiles)
-
-        if is_valid:
-            results['valid'] += 1
-            unique_smiles.add(corrected_smiles)
-            if corrected_smiles != processed_smiles:
-                results['corrected'] += 1
-
-            mol = Chem.MolFromSmiles(corrected_smiles)
-            if mol:
-                props = {
-                    'smiles': corrected_smiles,
-                    'MolWt': Descriptors.ExactMolWt(mol),
-                    'LogP': Descriptors.MolLogP(mol),
-                    'NumHDonors': Descriptors.NumHDonors(mol),
-                    'NumHAcceptors': Descriptors.NumHAcceptors(mol)
-                }
-
-                if corrected_smiles not in training_smiles_set:
-                    novel_smiles.add(corrected_smiles)
-                    results['novel'] += 1
-                    results['novel_properties'].append(props)
-                else:
-                    results['valid_properties'].append(props)
-        else:
-            results['invalid'] += 1
-            results['invalid_smiles'].append(smiles)
-
-    results['unique'] = len(unique_smiles)
-    return results
-
-# Function to visualize molecules
-def visualize_molecules(smiles_list: List[str], n: int = 5) -> Optional[Image.Image]:
-    valid_mols = []
-    for smiles in smiles_list:
-        smiles = smiles.strip().strip('<>').strip()
-        if not smiles:
-            continue
-        try:
-            mol = Chem.MolFromSmiles(smiles)
-            if mol is not None:
-                valid_mols.append(mol)
-                if len(valid_mols) == n:
-                    break
-        except Exception:
-            continue
-
-    if not valid_mols:
-        return None
-
-    try:
-        img = Draw.MolsToGridImage(
-            valid_mols,
-            molsPerRow=min(3, len(valid_mols)),
-            subImgSize=(200, 200),
-            legends=[f"Mol {i+1}" for i in range(len(valid_mols))]
-        )
-        return img
-    except Exception:
-        return None
-
-# GCN and GIN model definitions
-class GCN(torch.nn.Module):
-    """Graph Convolutional Network class with 3 convolutional layers and a linear layer"""
-
-    def __init__(self, dim_h):
-        """init method for GCN
-
-        Args:
-            dim_h (int): the dimension of hidden layers
-        """
-        super().__init__()
-        self.conv1 = GCNConv(11, dim_h)
-        self.conv2 = GCNConv(dim_h, dim_h)
-        self.conv3 = GCNConv(dim_h, dim_h)
-        self.lin = torch.nn.Linear(dim_h, 1)
-
-    def forward(self, data):
-        e = data.edge_index
-        x = data.x
-
-        x = self.conv1(x, e)
-        x = x.relu()
-        x = self.conv2(x, e)
-        x = x.relu()
-        x = self.conv3(x, e)
-        x = global_mean_pool(x, data.batch)
-
-        x = F.dropout(x, p=0.5, training=self.training)
-        x = self.lin(x)
-
-        return x
-
-class GIN(torch.nn.Module):
-    """Graph Isomorphism Network class with 3 GINConv layers and 2 linear layers"""
-
-    def __init__(self, dim_h):
-        """Initializing GIN class
-
-        Args:
-            dim_h (int): the dimension of hidden layers
-        """
-        super(GIN, self).__init__()
-        nn1 = Sequential(Linear(11, dim_h), BatchNorm1d(dim_h), ReLU(), Linear(dim_h, dim_h), ReLU())
-        self.conv1 = GINConv(nn1)
-        nn2 = Sequential(Linear(dim_h, dim_h), BatchNorm1d(dim_h), ReLU(), Linear(dim_h, dim_h), ReLU())
-        self.conv2 = GINConv(nn2)
-        nn3 = Sequential(Linear(dim_h, dim_h), BatchNorm1d(dim_h), ReLU(), Linear(dim_h, dim_h), ReLU())
-        self.conv3 = GINConv(nn3)
-        self.lin1 = Linear(dim_h, dim_h)
-        self.lin2 = Linear(dim_h, 1)
-
-    def forward(self, data):
-        x = data.x
-        edge_index = data.edge_index
-        batch = data.batch
-
-        # Node embeddings
-        h = self.conv1(x, edge_index)
-        h = h.relu()
-        h = self.conv2(h, edge_index)
-        h = h.relu()
-        h = self.conv3(h, edge_index)
-
-        # Graph-level readout
-        h = global_add_pool(h, batch)
-
-        h = self.lin1(h)
-        h = h.relu()
-        h = F.dropout(h, p=0.5, training=self.training)
-        h = self.lin2(h)
-
-        return h
-
-# Function to load GNN models
-@st.cache_resource
-def load_gnn_models(device):
-    # Load GCN model
-    gcn_model = GCN(dim_h=128)
-    gcn_model.load_state_dict(torch.load("GCN_model.pth", map_location=device))
-    gcn_model.to(device)
-    gcn_model.eval()
-
-    # Load GIN model
-    gin_model = GIN(dim_h=64)
-    gin_model.load_state_dict(torch.load("GIN_model.pth", map_location=device))
-    gin_model.to(device)
-    gin_model.eval()
-
-    return gcn_model, gin_model
-
-# Function to load normalization parameters
-@st.cache_resource
-def load_data_norm(device):
-    data_norm = torch.load('data_norm.pth', map_location=device)
-    data_mean = data_norm['mean']
-    data_std = data_norm['std']
-    return data_mean, data_std
-
-# Function to convert SMILES to Data object
-def smiles_to_data(smiles):
-    mol = Chem.MolFromSmiles(smiles)
-    if mol is None:
-        return None
-
-    atoms = mol.GetAtoms()
-    num_atoms = len(atoms)
-
-    atom_type_list = ['H', 'C', 'N', 'O', 'F']
-    hybridization_list = [Chem.rdchem.HybridizationType.SP, Chem.rdchem.HybridizationType.SP2, Chem.rdchem.HybridizationType.SP3]
-
-    x = []
-    for atom in atoms:
-        atom_type = atom.GetSymbol()
-        atom_type_feature = [int(atom_type == s) for s in atom_type_list]  # 5 features
-
-        # Atom degree (scalar between 0 and 4)
-        degree = atom.GetDegree()
-        degree_feature = [degree / 4]  # Normalize degree to [0,1]  # 1 feature
-
-        # Formal charge
-        formal_charge = atom.GetFormalCharge()
-        formal_charge_feature = [formal_charge / 4]  # Assume max formal charge is 4  # 1 feature
-
-        # Aromaticity
-        is_aromatic = atom.GetIsAromatic()
-        aromatic_feature = [int(is_aromatic)]  # 1 feature
-
-        # Hybridization
-        hybridization = atom.GetHybridization()
-        hybridization_feature = [int(hybridization == hyb) for hyb in hybridization_list]  # 3 features
-
-        # Total features: 5 + 1 +1 +1 +3 = 11
-        atom_feature = atom_type_feature + degree_feature + formal_charge_feature + aromatic_feature + hybridization_feature
-        x.append(atom_feature)
-
-    x = torch.tensor(x, dtype=torch.float)
-
-    # Build edge indices
-    edge_index = []
-    for bond in mol.GetBonds():
-        i = bond.GetBeginAtomIdx()
-        j = bond.GetEndAtomIdx()
-        edge_index.append([i, j])
-        edge_index.append([j, i])  # Since it's undirected
-
-    edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
-
-    # Build batch tensor (since batch size is 1)
-    batch = torch.zeros(num_atoms, dtype=torch.long)
-
-    # Build Data object
-    data = Data(x=x, edge_index=edge_index, batch=batch)
-
-    return data
-
-# Streamlit app
-def main():
-    st.set_page_config(
-        page_title="🧪 Molecule Generator and Property Predictor",
-        page_icon="🧪",
-        layout="wide",
-        initial_sidebar_state="expanded",
-    )
-
-    # Main Title and Description
-    st.title("🧪 Molecular Generation and Analysis using VAE and GNN")
-    st.markdown("""
-    This application allows you to generate novel molecular structures using a Variational Autoencoder (VAE) model trained on the QM9 dataset.
-    You can also predict molecular properties using pre-trained Graph Neural Network (GNN) models (GCN and GIN).
-    """)
-
-    # Initialize session state variables
-    if 'analysis' not in st.session_state:
-        st.session_state.analysis = None
-    if 'generated_smiles' not in st.session_state:
-        st.session_state.generated_smiles = []
-    if 'vae_generated' not in st.session_state:
-        st.session_state.vae_generated = False
-
-    # Sidebar configuration
-    st.sidebar.title("🔧 Configuration")
-    st.sidebar.markdown("Adjust the settings below to generate molecules or predict properties.")
-
-    # Load training data and canonicalize SMILES
-    @st.cache_data
-    def load_training_data():
-        df = pd.read_csv("qm9.csv")
-        smiles_list_raw = df['smiles'].tolist()
-        # Canonicalize SMILES for accurate comparison
-        smiles_list = [Chem.MolToSmiles(Chem.MolFromSmiles(s), isomericSmiles=True) for s in smiles_list_raw]
-        return set(smiles_list)
-
-    training_smiles_set = load_training_data()
-
-    # Device selection
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    # Model selection
-    st.sidebar.title("📌 Model Selection")
-    model_option = st.sidebar.selectbox("Choose a functionality", ("Generate Molecules (VAE)", "Predict Property (GNN)"))
-
-    if model_option == "Generate Molecules (VAE)":
-        # Number of samples
-        num_samples = st.sidebar.slider("Number of Molecules to Generate", min_value=5, max_value=500, value=50, step=5)
-
-        # Random seed
-        seed = st.sidebar.number_input("Random Seed", value=42, step=1)
-        torch.manual_seed(seed)
-        np.random.seed(seed)
-
-        if st.sidebar.button("🚀 Generate Molecules"):
-            with st.spinner("Generating molecules..."):
-                # Load VAE model
-                model, vocab = load_vae_model(device)
-                generated_smiles = generate_smiles_vae(model, vocab, num_samples=num_samples)
-                # Analyze molecules
-                analysis = analyze_molecules(generated_smiles, training_smiles_set)
-                # Store results in session state
-                st.session_state.generated_smiles = generated_smiles
-                st.session_state.analysis = analysis
-                st.session_state.vae_generated = True
-
-            # Display summary
-            st.success("✅ Molecule generation completed!")
-            st.subheader("Summary of Generated Molecules")
-            col1, col2, col3, col4 = st.columns(4)
-            col1.metric("Total Generated", analysis['total'])
-            col2.metric("Valid Molecules", f"{analysis['valid']} ({(analysis['valid']/analysis['total'])*100:.2f}%)")
-            col3.metric("Unique Molecules", f"{analysis['unique']} ({(analysis['unique']/analysis['total'])*100:.2f}%)")
-            col4.metric("Corrected Molecules", f"{analysis['corrected']} ({(analysis['corrected']/analysis['total'])*100:.2f}%)")
-
-            col1, col2 = st.columns(2)
-            col1.metric("Novel Molecules", f"{analysis['novel']} ({(analysis['novel']/analysis['total'])*100:.2f}%)")
-            col2.metric("Invalid Molecules", f"{analysis['invalid']} ({(analysis['invalid']/analysis['total'])*100:.2f}%)")
-
-            # Display properties
-            if analysis['valid_properties'] or analysis['novel_properties']:
-                st.subheader("Properties of Generated Molecules")
-
-                tab1, tab2 = st.tabs(["✅ Valid Molecules", "🌟 Novel Molecules"])
-                with tab1:
-                    if analysis['valid_properties']:
-                        df_valid = pd.DataFrame(analysis['valid_properties'])
-                        st.dataframe(df_valid)
-                        # Visualize valid molecules (limit to 9 for performance)
-                        st.subheader("Sample Valid Molecules")
-                        mol_image_valid = visualize_molecules([prop['smiles'] for prop in analysis['valid_properties']], n=9)
-                        if mol_image_valid:
-                            st.image(mol_image_valid)
-                        else:
-                            st.write("No valid molecules to display.")
-                    else:
-                        st.write("No valid molecules found.")
-
-                with tab2:
-                    if analysis['novel_properties']:
-                        df_novel = pd.DataFrame(analysis['novel_properties'])
-                        st.dataframe(df_novel)
-                        # Visualize novel molecules (limit to 9 for performance)
-                        st.subheader("Sample Novel Molecules")
-                        mol_image_novel = visualize_molecules([prop['smiles'] for prop in analysis['novel_properties']], n=9)
-                        if mol_image_novel:
-                            st.image(mol_image_novel)
-                        else:
-                            st.write("No novel molecules to display.")
-                    else:
-                        st.write("No novel molecules found.")
-
-                # Property distributions
-                st.subheader("Property Distributions")
-                fig, axs = plt.subplots(2, 2, figsize=(14, 10))
-                if analysis['valid_properties']:
-                    sns.histplot(df_valid['MolWt'], bins=20, ax=axs[0, 0], kde=True, color='skyblue', label='Valid')
-                if analysis['novel_properties']:
-                    sns.histplot(df_novel['MolWt'], bins=20, ax=axs[0, 0], kde=True, color='orange', label='Novel')
-                axs[0, 0].set_title('Molecular Weight Distribution')
-                axs[0, 0].legend()
-
-                if analysis['valid_properties']:
-                    sns.histplot(df_valid['LogP'], bins=20, ax=axs[0, 1], kde=True, color='skyblue', label='Valid')
-                if analysis['novel_properties']:
-                    sns.histplot(df_novel['LogP'], bins=20, ax=axs[0, 1], kde=True, color='orange', label='Novel')
-                axs[0, 1].set_title('LogP Distribution')
-                axs[0, 1].legend()
-
-                if analysis['valid_properties']:
-                    sns.histplot(df_valid['NumHDonors'], bins=range(0, max(df_valid['NumHDonors'].max(), 
-                                                                         df_novel['NumHDonors'].max()) + 2), 
-                                ax=axs[1, 0], kde=False, color='skyblue', label='Valid')
-                if analysis['novel_properties']:
-                    sns.histplot(df_novel['NumHDonors'], bins=range(0, max(df_valid['NumHDonors'].max(), 
-                                                                         df_novel['NumHDonors'].max()) + 2), 
-                                ax=axs[1, 0], kde=False, color='orange', label='Novel')
-                axs[1, 0].set_title('Number of H Donors')
-                axs[1, 0].legend()
-
-                if analysis['valid_properties']:
-                    sns.histplot(df_valid['NumHAcceptors'], bins=range(0, max(df_valid['NumHAcceptors'].max(), 
-                                                                            df_novel['NumHAcceptors'].max()) + 2), 
-                                ax=axs[1, 1], kde=False, color='skyblue', label='Valid')
-                if analysis['novel_properties']:
-                    sns.histplot(df_novel['NumHAcceptors'], bins=range(0, max(df_valid['NumHAcceptors'].max(), 
-                                                                            df_novel['NumHAcceptors'].max()) + 2), 
-                                ax=axs[1, 1], kde=False, color='orange', label='Novel')
-                axs[1, 1].set_title('Number of H Acceptors')
-                axs[1, 1].legend()
-
-                plt.tight_layout()
-                st.pyplot(fig)
-
-                # Download options
-                csv_valid = df_valid.to_csv(index=False).encode('utf-8')
-                csv_novel = df_novel.to_csv(index=False).encode('utf-8')
-                col1, col2 = st.columns(2)
-                with col1:
-                    st.download_button(
-                        label="💾 Download Valid Molecules CSV",
-                        data=csv_valid,
-                        file_name='valid_molecules.csv',
-                        mime='text/csv'
-                    )
-                with col2:
-                    st.download_button(
-                        label="💾 Download Novel Molecules CSV",
-                        data=csv_novel,
-                        file_name='novel_molecules.csv',
-                        mime='text/csv'
-                    )
-            else:
-                st.warning("No valid or novel molecules were generated.")
-
-    elif model_option == "Predict Property (GNN)":
-        # Load GNN models
-        with st.spinner("Loading GNN models..."):
-            gcn_model, gin_model = load_gnn_models(device)
-            # Load normalization parameters
-            data_mean, data_std = load_data_norm(device)
-
-        # GNN Model selection
-        gnn_model_option = st.sidebar.selectbox("Choose a GNN model", ("GCN", "GIN"))
-
-        st.subheader("🔍 Predict Molecular Property using GNN")
-        st.markdown("""
-        Input a SMILES string to predict the dipole moment using the selected GNN model.
-        """)
-
-        # User inputs a SMILES string
-        user_smiles = st.text_input("Enter a SMILES string for property prediction:", "")
-
-        if user_smiles:
-            data = smiles_to_data(user_smiles)
-            if data:
-                data = data.to(device)
-                if gnn_model_option == "GCN":
-                    prediction = gcn_model(data)
-                    prediction = prediction.item()
-                elif gnn_model_option == "GIN":
-                    prediction = gin_model(data)
-                    prediction = prediction.item()
-                # Unnormalize the prediction
-                prediction = prediction * data_std.item() + data_mean.item()
-                st.success(f"**Predicted Dipole Moment ({gnn_model_option}):** {prediction:.4f}")
-                # Display molecule
-                mol = Chem.MolFromSmiles(user_smiles)
-                if mol:
-                    st.subheader("Molecular Structure")
-                    st.image(Draw.MolToImage(mol, size=(300, 300)))
-            else:
-                st.error("❌ Invalid SMILES string.")
-
-        st.markdown("---")
-        st.markdown("### Or select a molecule from the generated molecules (if any).")
-
-        # Check if molecules have been generated
-        if st.session_state.vae_generated and st.session_state.analysis is not None:
-            # Combine valid and novel properties
-            all_properties = st.session_state.analysis['valid_properties'] + st.session_state.analysis['novel_properties']
-            if all_properties:
-                smiles_options = [prop['smiles'] for prop in all_properties]
-                selected_smiles = st.selectbox("Select a molecule", smiles_options)
-                if selected_smiles:
-                    data = smiles_to_data(selected_smiles)
-                    if data:
-                        data = data.to(device)
-                        if gnn_model_option == "GCN":
-                            prediction = gcn_model(data)
-                            prediction = prediction.item()
-                        elif gnn_model_option == "GIN":
-                            prediction = gin_model(data)
-                            prediction = prediction.item()
-                        # Unnormalize the prediction
-                        prediction = prediction * data_std.item() + data_mean.item()
-                        st.success(f"**Predicted Dipole Moment ({gnn_model_option}):** {prediction:.4f}")
-                        # Display molecule
-                        mol = Chem.MolFromSmiles(selected_smiles)
-                        if mol:
-                            st.subheader("Molecular Structure")
-                            st.image(Draw.MolToImage(mol, size=(300, 300)))
-                    else:
-                        st.error("❌ Invalid SMILES string.")
-            else:
-                st.info("🔍 No valid or novel molecules available.")
-        else:
-            st.info("🔍 No generated molecules available. Generate molecules using the VAE first.")
-
-    # About section
-    st.sidebar.title("ℹ️ About")
-    st.sidebar.info("""
-    **Molecule Generator and Property Predictor App**
-
-    This app uses a Variational Autoencoder (VAE) model and Graph Neural Networks (GNNs) to generate novel molecular structures and predict molecular properties.
-
-    - **Developed by**: Arjun, Kaustubh, and Nachiket
-    - **Hugging Face Repository**: [Your Hugging Face Repository](https://huggingface.co/YourRepositoryLink)
-    """)
-
-    # Hide Streamlit footer and header
-    hide_streamlit_style = """
-    <style>
-    footer {visibility: hidden;}
-    header {visibility: hidden;}
-    </style>
-    """
-    st.markdown(hide_streamlit_style, unsafe_allow_html=True)
-
-# Run the app
-if __name__ == "__main__":
-    main()
+# app.py
+
+import streamlit as st
+import torch
+import torch.nn as nn
+import pickle
+import numpy as np
+import pandas as pd
+from typing import List, Dict, Tuple, Optional
+
+# RDKit for molecule handling
+from rdkit import Chem
+from rdkit.Chem import Draw, Descriptors
+from rdkit import RDLogger
+RDLogger.DisableLog('rdApp.*')
+
+# Visualization libraries
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+# For generating images in Streamlit
+from PIL import Image
+
+# Suppress warnings in RDKit
+import warnings
+warnings.filterwarnings('ignore')
+
+# Set Seaborn style
+sns.set_style('whitegrid')
+
+# Additional imports for GNN
+import torch.nn.functional as F
+from torch.nn import Linear, Sequential, BatchNorm1d, ReLU
+
+from torch_geometric.data import Data
+from torch_geometric.nn import GCNConv, GINConv
+from torch_geometric.nn import global_mean_pool, global_add_pool
+
+# Function to load the VAE model
+@st.cache_resource
+def load_vae_model(device):
+    # Load the vocabulary
+    with open('vae_vocab.pkl', 'rb') as f:
+        vocab = pickle.load(f)
+    vocab_size = len(vocab)
+    
+    # Initialize the model with the same parameters
+    hidden_dim = 256  # Ensure this matches your trained model
+    latent_dim = 64   # Ensure this matches your trained model
+    
+    # Define the VAE class (same as in your training script)
+    class VAE(nn.Module):
+        def __init__(self, vocab_size: int, hidden_dim: int, latent_dim: int):
+            super(VAE, self).__init__()
+            self.vocab_size = vocab_size
+            self.hidden_dim = hidden_dim
+            self.latent_dim = latent_dim
+
+            self.encoder = nn.GRU(vocab_size, hidden_dim, batch_first=True)
+            self.fc_mu = nn.Linear(hidden_dim, latent_dim)
+            self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
+
+            self.decoder = nn.GRU(vocab_size + latent_dim, hidden_dim, batch_first=True)
+            self.fc_output = nn.Linear(hidden_dim, vocab_size)
+        
+        def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+            _, h = self.encoder(x)
+            h = h.squeeze(0)
+            return self.fc_mu(h), self.fc_logvar(h)
+
+        def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+            std = torch.exp(0.5 * logvar)
+            eps = torch.randn_like(std)
+            return mu + eps * std
+
+        def decode(self, z: torch.Tensor, max_length: int) -> torch.Tensor:
+            batch_size = z.size(0)
+            h = torch.zeros(1, batch_size, self.hidden_dim).to(z.device)
+            x = torch.zeros(batch_size, 1, self.vocab_size).to(z.device)
+            x[:, 0, vocab['<']] = 1  # Start token
+            outputs = []
+
+            for _ in range(max_length):
+                z_input = z.unsqueeze(1)
+                decoder_input = torch.cat([x, z_input], dim=2)
+                output, h = self.decoder(decoder_input, h)
+                output = self.fc_output(output)
+                outputs.append(output)
+                x = torch.softmax(output, dim=-1)
+
+            return torch.cat(outputs, dim=1)
+
+        def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            mu, logvar = self.encode(x)
+            z = self.reparameterize(mu, logvar)
+            return self.decode(z, x.size(1)), mu, logvar
+
+    model = VAE(vocab_size, hidden_dim, latent_dim)
+    model.load_state_dict(torch.load('vae_model.pth', map_location=device))
+    model.to(device)
+    model.eval()
+    return model, vocab
+
+# Function to generate molecules using VAE
+def generate_smiles_vae(model, vocab, num_samples=10, max_length=100):
+    model.eval()
+    inv_vocab = {v: k for k, v in vocab.items()}
+    generated_smiles = []
+    device = next(model.parameters()).device
+
+    with torch.no_grad():
+        for _ in range(num_samples):
+            z = torch.randn(1, model.latent_dim).to(device)
+            x = torch.zeros(1, 1, model.vocab_size).to(device)
+            x[0, 0, vocab['<']] = 1
+            h = torch.zeros(1, 1, model.hidden_dim).to(device)
+
+            smiles = ''
+            for _ in range(max_length):
+                z_input = z.unsqueeze(1)
+                decoder_input = torch.cat([x, z_input], dim=2)
+                output, h = model.decoder(decoder_input, h)
+                output = model.fc_output(output)
+
+                probs = torch.softmax(output.squeeze(0), dim=-1)
+                next_char = torch.multinomial(probs, 1).item()
+
+                if next_char == vocab['>']:
+                    break
+
+                smiles += inv_vocab.get(next_char, '')
+                x = torch.zeros(1, 1, model.vocab_size).to(device)
+                x[0, 0, next_char] = 1
+
+            generated_smiles.append(smiles)
+
+    return generated_smiles
+
+# Function to post-process and validate SMILES strings
+def enhanced_post_process_smiles(smiles: str) -> str:
+    smiles = smiles.replace('<', '').replace('>', '')
+    allowed_chars = set('CNOPSFIBrClcnops()[]=@+-#0123456789')
+    smiles = ''.join(c for c in smiles if c in allowed_chars)
+
+    # Balance parentheses
+    open_count = smiles.count('(')
+    close_count = smiles.count(')')
+    if open_count > close_count:
+        smiles += ')' * (open_count - close_count)
+    elif close_count > open_count:
+        smiles = '(' * (close_count - open_count) + smiles
+
+    # Replace invalid double bonds
+    smiles = smiles.replace('==', '=')
+
+    # Attempt to close unclosed rings
+    for i in range(1, 10):
+        if smiles.count(str(i)) % 2 != 0:
+            smiles += str(i)
+
+    return smiles
+
+def validate_and_correct_smiles(smiles: str) -> Tuple[bool, str]:
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is not None:
+        try:
+            Chem.SanitizeMol(mol)
+            return True, Chem.MolToSmiles(mol, isomericSmiles=True)
+        except:
+            pass
+    return False, smiles
+
+# Function to analyze molecules
+def analyze_molecules(smiles_list: List[str], training_smiles_set: set) -> Dict:
+    results = {
+        'total': len(smiles_list),
+        'valid': 0,
+        'invalid': 0,
+        'unique': 0,
+        'corrected': 0,
+        'novel': 0,
+        'valid_properties': [],
+        'novel_properties': [],
+        'invalid_smiles': []
+    }
+
+    unique_smiles = set()
+    novel_smiles = set()
+
+    for smiles in smiles_list:
+        processed_smiles = enhanced_post_process_smiles(smiles)
+        is_valid, corrected_smiles = validate_and_correct_smiles(processed_smiles)
+
+        if is_valid:
+            results['valid'] += 1
+            unique_smiles.add(corrected_smiles)
+            if corrected_smiles != processed_smiles:
+                results['corrected'] += 1
+
+            mol = Chem.MolFromSmiles(corrected_smiles)
+            if mol:
+                props = {
+                    'smiles': corrected_smiles,
+                    'MolWt': Descriptors.ExactMolWt(mol),
+                    'LogP': Descriptors.MolLogP(mol),
+                    'NumHDonors': Descriptors.NumHDonors(mol),
+                    'NumHAcceptors': Descriptors.NumHAcceptors(mol)
+                }
+
+                if corrected_smiles not in training_smiles_set:
+                    novel_smiles.add(corrected_smiles)
+                    results['novel'] += 1
+                    results['novel_properties'].append(props)
+                else:
+                    results['valid_properties'].append(props)
+        else:
+            results['invalid'] += 1
+            results['invalid_smiles'].append(smiles)
+
+    results['unique'] = len(unique_smiles)
+    return results
+
+# Function to visualize molecules
+def visualize_molecules(smiles_list: List[str], n: int = 5) -> Optional[Image.Image]:
+    valid_mols = []
+    for smiles in smiles_list:
+        smiles = smiles.strip().strip('<>').strip()
+        if not smiles:
+            continue
+        try:
+            mol = Chem.MolFromSmiles(smiles)
+            if mol is not None:
+                valid_mols.append(mol)
+                if len(valid_mols) == n:
+                    break
+        except Exception:
+            continue
+
+    if not valid_mols:
+        return None
+
+    try:
+        img = Draw.MolsToGridImage(
+            valid_mols,
+            molsPerRow=min(3, len(valid_mols)),
+            subImgSize=(200, 200),
+            legends=[f"Mol {i+1}" for i in range(len(valid_mols))]
+        )
+        return img
+    except Exception:
+        return None
+
+# GCN and GIN model definitions
+class GCN(torch.nn.Module):
+    """Graph Convolutional Network class with 3 convolutional layers and a linear layer"""
+
+    def __init__(self, dim_h):
+        """init method for GCN
+
+        Args:
+            dim_h (int): the dimension of hidden layers
+        """
+        super().__init__()
+        self.conv1 = GCNConv(11, dim_h)
+        self.conv2 = GCNConv(dim_h, dim_h)
+        self.conv3 = GCNConv(dim_h, dim_h)
+        self.lin = torch.nn.Linear(dim_h, 1)
+
+    def forward(self, data):
+        e = data.edge_index
+        x = data.x
+
+        x = self.conv1(x, e)
+        x = x.relu()
+        x = self.conv2(x, e)
+        x = x.relu()
+        x = self.conv3(x, e)
+        x = global_mean_pool(x, data.batch)
+
+        x = F.dropout(x, p=0.5, training=self.training)
+        x = self.lin(x)
+
+        return x
+
+class GIN(torch.nn.Module):
+    """Graph Isomorphism Network class with 3 GINConv layers and 2 linear layers"""
+
+    def __init__(self, dim_h):
+        """Initializing GIN class
+
+        Args:
+            dim_h (int): the dimension of hidden layers
+        """
+        super(GIN, self).__init__()
+        nn1 = Sequential(Linear(11, dim_h), BatchNorm1d(dim_h), ReLU(), Linear(dim_h, dim_h), ReLU())
+        self.conv1 = GINConv(nn1)
+        nn2 = Sequential(Linear(dim_h, dim_h), BatchNorm1d(dim_h), ReLU(), Linear(dim_h, dim_h), ReLU())
+        self.conv2 = GINConv(nn2)
+        nn3 = Sequential(Linear(dim_h, dim_h), BatchNorm1d(dim_h), ReLU(), Linear(dim_h, dim_h), ReLU())
+        self.conv3 = GINConv(nn3)
+        self.lin1 = Linear(dim_h, dim_h)
+        self.lin2 = Linear(dim_h, 1)
+
+    def forward(self, data):
+        x = data.x
+        edge_index = data.edge_index
+        batch = data.batch
+
+        # Node embeddings
+        h = self.conv1(x, edge_index)
+        h = h.relu()
+        h = self.conv2(h, edge_index)
+        h = h.relu()
+        h = self.conv3(h, edge_index)
+
+        # Graph-level readout
+        h = global_add_pool(h, batch)
+
+        h = self.lin1(h)
+        h = h.relu()
+        h = F.dropout(h, p=0.5, training=self.training)
+        h = self.lin2(h)
+
+        return h
+
+# Function to load GNN models
+@st.cache_resource
+def load_gnn_models(device):
+    # Load GCN model
+    gcn_model = GCN(dim_h=128)
+    gcn_model.load_state_dict(torch.load("GCN_model.pth", map_location=device))
+    gcn_model.to(device)
+    gcn_model.eval()
+
+    # Load GIN model
+    gin_model = GIN(dim_h=64)
+    gin_model.load_state_dict(torch.load("GIN_model.pth", map_location=device))
+    gin_model.to(device)
+    gin_model.eval()
+
+    return gcn_model, gin_model
+
+# Function to load normalization parameters
+@st.cache_resource
+def load_data_norm(device):
+    data_norm = torch.load('data_norm.pth', map_location=device)
+    data_mean = data_norm['mean']
+    data_std = data_norm['std']
+    return data_mean, data_std
+
+# Function to convert SMILES to Data object
+def smiles_to_data(smiles):
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return None
+
+    atoms = mol.GetAtoms()
+    num_atoms = len(atoms)
+
+    atom_type_list = ['H', 'C', 'N', 'O', 'F']
+    hybridization_list = [Chem.rdchem.HybridizationType.SP, Chem.rdchem.HybridizationType.SP2, Chem.rdchem.HybridizationType.SP3]
+
+    x = []
+    for atom in atoms:
+        atom_type = atom.GetSymbol()
+        atom_type_feature = [int(atom_type == s) for s in atom_type_list]  # 5 features
+
+        # Atom degree (scalar between 0 and 4)
+        degree = atom.GetDegree()
+        degree_feature = [degree / 4]  # Normalize degree to [0,1]  # 1 feature
+
+        # Formal charge
+        formal_charge = atom.GetFormalCharge()
+        formal_charge_feature = [formal_charge / 4]  # Assume max formal charge is 4  # 1 feature
+
+        # Aromaticity
+        is_aromatic = atom.GetIsAromatic()
+        aromatic_feature = [int(is_aromatic)]  # 1 feature
+
+        # Hybridization
+        hybridization = atom.GetHybridization()
+        hybridization_feature = [int(hybridization == hyb) for hyb in hybridization_list]  # 3 features
+
+        # Total features: 5 + 1 +1 +1 +3 = 11
+        atom_feature = atom_type_feature + degree_feature + formal_charge_feature + aromatic_feature + hybridization_feature
+        x.append(atom_feature)
+
+    x = torch.tensor(x, dtype=torch.float)
+
+    # Build edge indices
+    edge_index = []
+    for bond in mol.GetBonds():
+        i = bond.GetBeginAtomIdx()
+        j = bond.GetEndAtomIdx()
+        edge_index.append([i, j])
+        edge_index.append([j, i])  # Since it's undirected
+
+    edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
+
+    # Build batch tensor (since batch size is 1)
+    batch = torch.zeros(num_atoms, dtype=torch.long)
+
+    # Build Data object
+    data = Data(x=x, edge_index=edge_index, batch=batch)
+
+    return data
+
+# Streamlit app
+def main():
+    st.set_page_config(
+        page_title="🧪 Molecule Generator and Property Predictor",
+        page_icon="🧪",
+        layout="wide",
+        initial_sidebar_state="expanded",
+    )
+
+    # Main Title and Description
+    st.title("🧪 Molecular Generation and Analysis using VAE and GNN")
+    st.markdown("""
+    This application allows you to generate novel molecular structures using a Variational Autoencoder (VAE) model trained on the QM9 dataset.
+    You can also predict molecular properties using pre-trained Graph Neural Network (GNN) models (GCN and GIN).
+    """)
+
+    # Initialize session state variables
+    if 'analysis' not in st.session_state:
+        st.session_state.analysis = None
+    if 'generated_smiles' not in st.session_state:
+        st.session_state.generated_smiles = []
+    if 'vae_generated' not in st.session_state:
+        st.session_state.vae_generated = False
+
+    # Sidebar configuration
+    st.sidebar.title("🔧 Configuration")
+    st.sidebar.markdown("Adjust the settings below to generate molecules or predict properties.")
+
+    # Load training data and canonicalize SMILES
+    @st.cache_data
+    def load_training_data():
+        df = pd.read_csv("qm9.csv")
+        smiles_list_raw = df['smiles'].tolist()
+        # Canonicalize SMILES for accurate comparison
+        smiles_list = [Chem.MolToSmiles(Chem.MolFromSmiles(s), isomericSmiles=True) for s in smiles_list_raw]
+        return set(smiles_list)
+
+    training_smiles_set = load_training_data()
+
+    # Device selection
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Model selection
+    st.sidebar.title("📌 Model Selection")
+    model_option = st.sidebar.selectbox("Choose a functionality", ("Generate Molecules (VAE)", "Predict Property (GNN)"))
+
+    if model_option == "Generate Molecules (VAE)":
+        # Number of samples
+        num_samples = st.sidebar.slider("Number of Molecules to Generate", min_value=5, max_value=500, value=50, step=5)
+
+        # Random seed
+        seed = st.sidebar.number_input("Random Seed", value=42, step=1)
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+
+        if st.sidebar.button("🚀 Generate Molecules"):
+            with st.spinner("Generating molecules..."):
+                # Load VAE model
+                model, vocab = load_vae_model(device)
+                generated_smiles = generate_smiles_vae(model, vocab, num_samples=num_samples)
+                # Analyze molecules
+                analysis = analyze_molecules(generated_smiles, training_smiles_set)
+                # Store results in session state
+                st.session_state.generated_smiles = generated_smiles
+                st.session_state.analysis = analysis
+                st.session_state.vae_generated = True
+
+            # Display summary
+            st.success("✅ Molecule generation completed!")
+            st.subheader("Summary of Generated Molecules")
+            col1, col2, col3, col4 = st.columns(4)
+            col1.metric("Total Generated", analysis['total'])
+            col2.metric("Valid Molecules", f"{analysis['valid']} ({(analysis['valid']/analysis['total'])*100:.2f}%)")
+            col3.metric("Unique Molecules", f"{analysis['unique']} ({(analysis['unique']/analysis['total'])*100:.2f}%)")
+            col4.metric("Corrected Molecules", f"{analysis['corrected']} ({(analysis['corrected']/analysis['total'])*100:.2f}%)")
+
+            col1, col2 = st.columns(2)
+            col1.metric("Novel Molecules", f"{analysis['novel']} ({(analysis['novel']/analysis['total'])*100:.2f}%)")
+            col2.metric("Invalid Molecules", f"{analysis['invalid']} ({(analysis['invalid']/analysis['total'])*100:.2f}%)")
+
+            # Display properties
+            if analysis['valid_properties'] or analysis['novel_properties']:
+                st.subheader("Properties of Generated Molecules")
+
+                tab1, tab2 = st.tabs(["✅ Valid Molecules", "🌟 Novel Molecules"])
+                with tab1:
+                    if analysis['valid_properties']:
+                        df_valid = pd.DataFrame(analysis['valid_properties'])
+                        st.dataframe(df_valid)
+                        # Visualize valid molecules (limit to 9 for performance)
+                        st.subheader("Sample Valid Molecules")
+                        mol_image_valid = visualize_molecules([prop['smiles'] for prop in analysis['valid_properties']], n=9)
+                        if mol_image_valid:
+                            st.image(mol_image_valid)
+                        else:
+                            st.write("No valid molecules to display.")
+                    else:
+                        st.write("No valid molecules found.")
+
+                with tab2:
+                    if analysis['novel_properties']:
+                        df_novel = pd.DataFrame(analysis['novel_properties'])
+                        st.dataframe(df_novel)
+                        # Visualize novel molecules (limit to 9 for performance)
+                        st.subheader("Sample Novel Molecules")
+                        mol_image_novel = visualize_molecules([prop['smiles'] for prop in analysis['novel_properties']], n=9)
+                        if mol_image_novel:
+                            st.image(mol_image_novel)
+                        else:
+                            st.write("No novel molecules to display.")
+                    else:
+                        st.write("No novel molecules found.")
+
+                # Property distributions
+                st.subheader("Property Distributions")
+                fig, axs = plt.subplots(2, 2, figsize=(14, 10))
+                if analysis['valid_properties']:
+                    sns.histplot(df_valid['MolWt'], bins=20, ax=axs[0, 0], kde=True, color='skyblue', label='Valid')
+                if analysis['novel_properties']:
+                    sns.histplot(df_novel['MolWt'], bins=20, ax=axs[0, 0], kde=True, color='orange', label='Novel')
+                axs[0, 0].set_title('Molecular Weight Distribution')
+                axs[0, 0].legend()
+
+                if analysis['valid_properties']:
+                    sns.histplot(df_valid['LogP'], bins=20, ax=axs[0, 1], kde=True, color='skyblue', label='Valid')
+                if analysis['novel_properties']:
+                    sns.histplot(df_novel['LogP'], bins=20, ax=axs[0, 1], kde=True, color='orange', label='Novel')
+                axs[0, 1].set_title('LogP Distribution')
+                axs[0, 1].legend()
+
+                if analysis['valid_properties']:
+                    sns.histplot(df_valid['NumHDonors'], bins=range(0, max(df_valid['NumHDonors'].max(), 
+                                                                         df_novel['NumHDonors'].max()) + 2), 
+                                ax=axs[1, 0], kde=False, color='skyblue', label='Valid')
+                if analysis['novel_properties']:
+                    sns.histplot(df_novel['NumHDonors'], bins=range(0, max(df_valid['NumHDonors'].max(), 
+                                                                         df_novel['NumHDonors'].max()) + 2), 
+                                ax=axs[1, 0], kde=False, color='orange', label='Novel')
+                axs[1, 0].set_title('Number of H Donors')
+                axs[1, 0].legend()
+
+                if analysis['valid_properties']:
+                    sns.histplot(df_valid['NumHAcceptors'], bins=range(0, max(df_valid['NumHAcceptors'].max(), 
+                                                                            df_novel['NumHAcceptors'].max()) + 2), 
+                                ax=axs[1, 1], kde=False, color='skyblue', label='Valid')
+                if analysis['novel_properties']:
+                    sns.histplot(df_novel['NumHAcceptors'], bins=range(0, max(df_valid['NumHAcceptors'].max(), 
+                                                                            df_novel['NumHAcceptors'].max()) + 2), 
+                                ax=axs[1, 1], kde=False, color='orange', label='Novel')
+                axs[1, 1].set_title('Number of H Acceptors')
+                axs[1, 1].legend()
+
+                plt.tight_layout()
+                st.pyplot(fig)
+
+                # Download options
+                csv_valid = df_valid.to_csv(index=False).encode('utf-8')
+                csv_novel = df_novel.to_csv(index=False).encode('utf-8')
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.download_button(
+                        label="💾 Download Valid Molecules CSV",
+                        data=csv_valid,
+                        file_name='valid_molecules.csv',
+                        mime='text/csv'
+                    )
+                with col2:
+                    st.download_button(
+                        label="💾 Download Novel Molecules CSV",
+                        data=csv_novel,
+                        file_name='novel_molecules.csv',
+                        mime='text/csv'
+                    )
+            else:
+                st.warning("No valid or novel molecules were generated.")
+
+    elif model_option == "Predict Property (GNN)":
+        # Load GNN models
+        with st.spinner("Loading GNN models..."):
+            gcn_model, gin_model = load_gnn_models(device)
+            # Load normalization parameters
+            data_mean, data_std = load_data_norm(device)
+
+        # GNN Model selection
+        gnn_model_option = st.sidebar.selectbox("Choose a GNN model", ("GCN", "GIN"))
+
+        st.subheader("🔍 Predict Molecular Property using GNN")
+        st.markdown("""
+        Input a SMILES string to predict the dipole moment using the selected GNN model.
+        """)
+
+        # User inputs a SMILES string
+        user_smiles = st.text_input("Enter a SMILES string for property prediction:", "")
+
+        if user_smiles:
+            data = smiles_to_data(user_smiles)
+            if data:
+                data = data.to(device)
+                if gnn_model_option == "GCN":
+                    prediction = gcn_model(data)
+                    prediction = prediction.item()
+                elif gnn_model_option == "GIN":
+                    prediction = gin_model(data)
+                    prediction = prediction.item()
+                # Unnormalize the prediction
+                prediction = prediction * data_std.item() + data_mean.item()
+                st.success(f"**Predicted Dipole Moment ({gnn_model_option}):** {prediction:.4f}")
+                # Display molecule
+                mol = Chem.MolFromSmiles(user_smiles)
+                if mol:
+                    st.subheader("Molecular Structure")
+                    st.image(Draw.MolToImage(mol, size=(300, 300)))
+            else:
+                st.error("❌ Invalid SMILES string.")
+
+        st.markdown("---")
+        st.markdown("### Or select a molecule from the generated molecules (if any).")
+
+        # Check if molecules have been generated
+        if st.session_state.vae_generated and st.session_state.analysis is not None:
+            # Combine valid and novel properties
+            all_properties = st.session_state.analysis['valid_properties'] + st.session_state.analysis['novel_properties']
+            if all_properties:
+                smiles_options = [prop['smiles'] for prop in all_properties]
+                selected_smiles = st.selectbox("Select a molecule", smiles_options)
+                if selected_smiles:
+                    data = smiles_to_data(selected_smiles)
+                    if data:
+                        data = data.to(device)
+                        if gnn_model_option == "GCN":
+                            prediction = gcn_model(data)
+                            prediction = prediction.item()
+                        elif gnn_model_option == "GIN":
+                            prediction = gin_model(data)
+                            prediction = prediction.item()
+                        # Unnormalize the prediction
+                        prediction = prediction * data_std.item() + data_mean.item()
+                        st.success(f"**Predicted Dipole Moment ({gnn_model_option}):** {prediction:.4f}")
+                        # Display molecule
+                        mol = Chem.MolFromSmiles(selected_smiles)
+                        if mol:
+                            st.subheader("Molecular Structure")
+                            st.image(Draw.MolToImage(mol, size=(300, 300)))
+                    else:
+                        st.error("❌ Invalid SMILES string.")
+            else:
+                st.info("🔍 No valid or novel molecules available.")
+        else:
+            st.info("🔍 No generated molecules available. Generate molecules using the VAE first.")
+
+    # About section
+    st.sidebar.title("ℹ️ About")
+    st.sidebar.info("""
+    **Molecule Generator and Property Predictor App**
+
+    This app uses a Variational Autoencoder (VAE) model and Graph Neural Networks (GNNs) to generate novel molecular structures and predict molecular properties.
+
+    - **Developed by**: Arjun, Kaustubh, and Nachiket
+    - **Hugging Face Repository**: https://huggingface.co/spaces/Raykarr/SMILES_Generation_and_Prediction
+    """)
+
+    # Hide Streamlit footer and header
+    hide_streamlit_style = """
+    <style>
+    footer {visibility: hidden;}
+    header {visibility: hidden;}
+    </style>
+    """
+    st.markdown(hide_streamlit_style, unsafe_allow_html=True)
+
+# Run the app
+if __name__ == "__main__":
+    main()