Upload app.py with huggingface_hub

app.py

@@ -1,136 +1,203 @@
 import spaces
+import gradio as gr
 
-# Standard Libraries
-import os
-import io
-import csv
-import json
-import glob
+import torch
+import numpy as np
+import pandas as pd
 import random
+import io
+import imageio
+import os
 import tempfile
 import atexit
+import glob
+import csv
 from datetime import datetime
+import json
 
-# Third-Party Libraries
-import numpy as np
-import pandas as pd
-import torch
-import imageio
 from rdkit import Chem
 from rdkit.Chem import Draw
-import gradio as gr
 
-# Local Modules
 from evaluator import Evaluator
-from loader import load_graph_decoder
-
-# --------------------------- Configuration Constants --------------------------- #
-
-DATA_DIR = 'data'
-EVALUATORS_DIR = 'evaluators'
-FLAGGED_FOLDER = "flagged"
-KNOWN_LABELS_FILE = os.path.join(DATA_DIR, 'known_labels.csv')
-KNOWN_SMILES_FILE = os.path.join(DATA_DIR, 'known_polymers.csv')
-
-ALL_PROPERTIES = ['CH4', 'CO2', 'H2', 'N2', 'O2']
-MODEL_NAME_MAPPING = {
-    "model_all": "Graph DiT (trained on labeled + unlabeled)",
-    "model_labeled": "Graph DiT (trained on labeled)"
-}
-
-GIF_TEMP_PREFIX = "polymer_gifs_"
-
-# --------------------------- Data Loading --------------------------- #
-
-def load_known_data():
-    """Load known labels and SMILES data from CSV files."""
-    try:
-        known_labels = pd.read_csv(KNOWN_LABELS_FILE)
-        known_smiles = pd.read_csv(KNOWN_SMILES_FILE)
-        return known_labels, known_smiles
-    except Exception as e:
-        raise FileNotFoundError(f"Error loading data files: {e}")
-
-# Load data
-known_labels, known_smiles = load_known_data()
+# from loader import load_graph_decoder
 
-# --------------------------- Evaluator Setup --------------------------- #
-
-def initialize_evaluators(properties, evaluators_dir):
-    """Initialize evaluators for each property."""
-    evaluators = {}
-    for prop in properties:
-        evaluator_path = os.path.join(evaluators_dir, f'{prop}.joblib')
-        evaluators[prop] = Evaluator(evaluator_path, prop)
-    return evaluators
+### load model start
+from graph_decoder.diffusion_model import GraphDiT
+def count_parameters(model):
+    r"""
+    Returns the number of trainable parameters and number of all parameters in the model.
+    """
+    trainable_params, all_param = 0, 0
+    for param in model.parameters():
+        num_params = param.numel()
+        all_param += num_params
+        if param.requires_grad:
+            trainable_params += num_params
+
+    return trainable_params, all_param
+
+def load_graph_decoder(path='model_labeled'):
+    model_config_path = f"{path}/config.yaml"
+    data_info_path = f"{path}/data.meta.json"
+
+    model = GraphDiT(
+        model_config_path=model_config_path,
+        data_info_path=data_info_path,
+        model_dtype=torch.float32,
+    )
+    model.init_model(path)
+    model.disable_grads()
 
-evaluators = initialize_evaluators(ALL_PROPERTIES, EVALUATORS_DIR)
+    trainable_params, all_param = count_parameters(model)
+    param_stats = "Loaded Graph DiT from {} trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format(
+        path, trainable_params, all_param, 100 * trainable_params / all_param
+    )
+    print(param_stats)
+    return model
+### load model end
 
-# --------------------------- Property Ranges --------------------------- #
+# Load the CSV data
+known_labels = pd.read_csv('data/known_labels.csv')
+knwon_smiles = pd.read_csv('data/known_polymers.csv')
 
-def get_property_ranges(labels, properties):
-    """Get min and max values for each property."""
-    return {prop: (labels[prop].min(), labels[prop].max()) for prop in properties}
+all_properties = ['CH4', 'CO2', 'H2', 'N2', 'O2']
 
-property_ranges = get_property_ranges(known_labels, ALL_PROPERTIES)
+# Initialize evaluators
+evaluators = {prop: Evaluator(f'evaluators/{prop}.joblib', prop) for prop in all_properties}
 
-# --------------------------- Temporary Directory Setup --------------------------- #
+# Get min and max values for each property
+property_ranges = {prop: (known_labels[prop].min(), known_labels[prop].max()) for prop in all_properties}
 
-temp_dir = tempfile.mkdtemp(prefix=GIF_TEMP_PREFIX)
+# Create a temporary directory for GIFs
+temp_dir = tempfile.mkdtemp(prefix="polymer_gifs_")
 
 def cleanup_temp_files():
     """Clean up temporary GIF files on exit."""
-    try:
-        for file in glob.glob(os.path.join(temp_dir, "*.gif")):
+    for file in glob.glob(os.path.join(temp_dir, "*.gif")):
+        try:
             os.remove(file)
+        except Exception as e:
+            print(f"Error deleting {file}: {e}")
+    try:
         os.rmdir(temp_dir)
     except Exception as e:
-        print(f"Error during cleanup: {e}")
+        print(f"Error deleting temporary directory {temp_dir}: {e}")
 
+# Register the cleanup function to be called on exit
 atexit.register(cleanup_temp_files)
 
-# --------------------------- Utility Functions --------------------------- #
-
 def random_properties():
-    """Select a random set of properties from known labels."""
-    return known_labels[ALL_PROPERTIES].sample(1).values.tolist()[0]
+    return known_labels[all_properties].sample(1).values.tolist()[0]
 
 def load_model(model_choice):
-    """Load the graph decoder model based on the choice."""
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = load_graph_decoder(path=model_choice)
-    return model, device
+    return (model, device)
+
+# Create a flagged folder if it doesn't exist
+flagged_folder = "flagged"
+os.makedirs(flagged_folder, exist_ok=True)
 
 def save_interesting_log(smiles, properties, suggested_properties):
-    """Save interesting polymer data to a CSV log file."""
-    log_file = os.path.join(FLAGGED_FOLDER, "log.csv")
-    os.makedirs(FLAGGED_FOLDER, exist_ok=True)
+    """Save interesting polymer data to a CSV file."""
+    log_file = os.path.join(flagged_folder, "log.csv")
     file_exists = os.path.isfile(log_file)
+    
+    with open(log_file, 'a', newline='') as csvfile:
+        fieldnames = ['timestamp', 'smiles'] + all_properties + [f'suggested_{prop}' for prop in all_properties]
+        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+        
+        if not file_exists:
+            writer.writeheader()
+        
+        log_data = {
+            'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            'smiles': smiles,
+            **{prop: value for prop, value in zip(all_properties, properties)},
+            **{f'suggested_{prop}': value for prop, value in suggested_properties.items()}
+        }
+        writer.writerow(log_data)
 
-    fieldnames = ['timestamp', 'smiles'] + ALL_PROPERTIES + [f'suggested_{prop}' for prop in ALL_PROPERTIES]
+@spaces.GPU
+def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
+    print('in generate_graph')
+    model, device = model_state
+    
+    properties = [CH4, CO2, H2, N2, O2]
+    
+    def is_nan_like(x):
+        return x == 0 or x == '' or (isinstance(x, float) and np.isnan(x))
+    
+    properties = [None if is_nan_like(prop) else prop for prop in properties]
+    
+    nan_message = "The following gas properties were treated as NaN: "
+    nan_gases = [gas for gas, prop in zip(all_properties, properties) if prop is None]
+    nan_message += ", ".join(nan_gases) if nan_gases else "None"
 
-    try:
-        with open(log_file, 'a', newline='') as csvfile:
-            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
-            if not file_exists:
-                writer.writeheader()
-
-            log_data = {
-                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-                'smiles': smiles,
-                **{prop: value for prop, value in zip(ALL_PROPERTIES, properties)},
-                **{f'suggested_{prop}': value for prop, value in suggested_properties.items()}
-            }
-            writer.writerow(log_data)
-    except Exception as e:
-        print(f"Error saving log: {e}")
+    num_nodes = None if num_nodes == 0 else num_nodes
+    
+    for _ in range(repeating_time):
+        # try:
+        model.to(device)
+        generated_molecule, img_list = model.generate(properties, device=device, guide_scale=guidance_scale, num_nodes=num_nodes, number_chain_steps=num_chain_steps)
+
+        # Create GIF if img_list is available
+        gif_path = None
+        if img_list and len(img_list) > 0:
+            imgs = [np.array(pil_img) for pil_img in img_list]
+            imgs.extend([imgs[-1]] * 10)
+            gif_path = os.path.join(temp_dir, f"polymer_gen_{random.randint(0, 999999)}.gif")
+            imageio.mimsave(gif_path, imgs, format='GIF', fps=fps, loop=0)
+        
+        if generated_molecule is not None:
+            mol = Chem.MolFromSmiles(generated_molecule)
+            if mol is not None:
+                standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
+                is_novel = standardized_smiles not in knwon_smiles['SMILES'].values
+                novelty_status = "Novel (Not in Labeled Set)" if is_novel else "Not Novel (Exists in Labeled Set)"
+                img = Draw.MolToImage(mol)
+                
+                # Evaluate the generated molecule
+                suggested_properties = {}
+                for prop, evaluator in evaluators.items():
+                    suggested_properties[prop] = evaluator([standardized_smiles])[0]
+                
+                suggested_properties_text = "\n".join([f"**Suggested {prop}:** {value:.2f}" for prop, value in suggested_properties.items()])
+                
+                return (
+                    f"**Generated polymer SMILES:** `{standardized_smiles}`\n\n"
+                    f"**{nan_message}**\n\n"
+                    f"**{novelty_status}**\n\n"
+                    f"**Suggested Properties:**\n{suggested_properties_text}",
+                    img,
+                    gif_path,
+                    properties,  # Add this
+                    suggested_properties  # Add this
+                )
+        else:
+            return (
+                f"**Generation failed:** Could not generate a valid molecule.\n\n**{nan_message}**",
+                None,
+                gif_path,
+                properties,
+                None,
+            )
+        # except Exception as e:
+        #     print(f"Error in generation: {e}")
+        #     continue
+    
+    return f"**Generation failed:** Could not generate a valid molecule after {repeating_time} attempts.\n\n**{nan_message}**", None, None
 
-def is_nan_like(x):
-    """Check if a value should be treated as NaN."""
-    return x == 0 or x == '' or (isinstance(x, float) and np.isnan(x))
+def set_random_properties():
+    return random_properties()
+
+# Create a mapping of internal names to display names
+model_name_mapping = {
+    "model_all": "Graph DiT (trained on labeled + unlabeled)",
+    "model_labeled": "Graph DiT (trained on labeled)"
+}
 
 def numpy_to_python(obj):
-    """Convert NumPy objects to native Python types."""
     if isinstance(obj, np.integer):
         return int(obj)
     elif isinstance(obj, np.floating):
@@ -143,329 +210,175 @@ def numpy_to_python(obj):
         return {k: numpy_to_python(v) for k, v in obj.items()}
     else:
         return obj
+        
+def on_generate(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
+    print('in on_generate', on_generate)
+    result = generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps)
+    # Check if the generation was successful
+    if result[0].startswith("**Generated polymer SMILES:**"):
+        smiles = result[0].split("**Generated polymer SMILES:** `")[1].split("`")[0]
+        properties = json.dumps(numpy_to_python(result[3]))
+        suggested_properties = json.dumps(numpy_to_python(result[4]))
+        # Return the result with an enabled feedback button
+        return [*result[:3], smiles, properties, suggested_properties, gr.Button(interactive=True)]
+    else:
+        # Return the result with a disabled feedback button
+        return [*result[:3], "", "[]", "[]", gr.Button(interactive=False)]
 
-# --------------------------- Graph Generation Function --------------------------- #
-
-@spaces.GPU
-def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
-    """
-    Generate a polymer graph based on the input properties and model.
-    Returns generation results including SMILES, images, and properties.
-    """
-    print('Generating graph...')
-    model, device = model_state
-    properties = [CH4, CO2, H2, N2, O2]
+def process_feedback(checkbox_value, smiles, properties, suggested_properties):
+    if checkbox_value:
+        # Check if properties and suggested_properties are already Python objects
+        if isinstance(properties, str):
+            properties = json.loads(properties)
+        if isinstance(suggested_properties, str):
+            suggested_properties = json.loads(suggested_properties)
+        
+        save_interesting_log(smiles, properties, suggested_properties)
+        return gr.Textbox(value="Thank you for your feedback! This polymer has been saved to our interesting polymers log.", visible=True)
+    else:
+        return gr.Textbox(value="Thank you for your feedback!", visible=True)
 
-    # Handle NaN-like values
-    properties = [None if is_nan_like(prop) else prop for prop in properties]
+# ADD THIS FUNCTION
+def reset_feedback_button():
+    return gr.Button(interactive=False)
 
-    nan_gases = [gas for gas, prop in zip(ALL_PROPERTIES, properties) if prop is None]
-    nan_message = "The following gas properties were treated as NaN: " + (", ".join(nan_gases) if nan_gases else "None")
+# Create the Gradio interface using Blocks
+with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
+    # Navigation Bar
+    with gr.Row(elem_id="navbar"):
+        gr.Markdown("""
+        <div style="text-align: center;">
+            <h1>🔗🔬 Polymer Design with GraphDiT</h1>
+            <div style="display: flex; gap: 20px; justify-content: center; align-items: center; margin-top: 10px;">
+                <a href="https://github.com/liugangcode/Graph-DiT" target="_blank" style="display: flex; align-items: center; gap: 5px; text-decoration: none; color: inherit;">
+                    <img src="https://img.icons8.com/ios-glyphs/30/000000/github.png" alt="GitHub" />
+                    <span>View Code</span>
+                </a>
+                <a href="https://arxiv.org/abs/2401.13858" target="_blank" style="text-decoration: none; color: inherit;">
+                    📄 View Paper
+                </a>
+            </div>
+        </div>
+        """)
 
-    num_nodes = None if num_nodes == 0 else num_nodes
+    # Main Description
+    gr.Markdown("""
+    ## Introduction
 
-    for attempt in range(repeating_time):
-        try:
-            generated_molecule, img_list = model.generate(
-                properties,
-                device=device,
-                guide_scale=guidance_scale,
-                num_nodes=num_nodes,
-                number_chain_steps=num_chain_steps
-            )
+    Input the desired gas barrier properties for CH₄, CO₂, H₂, N₂, and O₂ to generate novel polymer structures. The results are visualized as molecular graphs and represented by SMILES strings if they are successfully generated. Note: Gas barrier values set to 0 will be treated as `NaN` (unconditionally). If the generation fails, please retry or increase the number of repetition attempts.
+    """)
 
-            gif_path = None
-            if img_list:
-                imgs = [np.array(pil_img) for pil_img in img_list]
-                imgs.extend([imgs[-1]] * 10)  # Extend the last image for GIF
-                gif_path = os.path.join(temp_dir, f"polymer_gen_{random.randint(0, 999999)}.gif")
-                imageio.mimsave(gif_path, imgs, format='GIF', fps=fps, loop=0)
-
-            if generated_molecule:
-                mol = Chem.MolFromSmiles(generated_molecule)
-                if mol:
-                    standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
-                    is_novel = standardized_smiles not in known_smiles['SMILES'].values
-                    novelty_status = "Novel (Not in Labeled Set)" if is_novel else "Not Novel (Exists in Labeled Set)"
-                    img = Draw.MolToImage(mol)
-
-                    # Evaluate the generated molecule
-                    suggested_properties = {prop: evaluator([standardized_smiles])[0] for prop, evaluator in evaluators.items()}
-
-                    suggested_properties_text = "\n".join([f"**Suggested {prop}:** {value:.2f}" for prop, value in suggested_properties.items()])
-
-                    return (
-                        f"**Generated polymer SMILES:** `{standardized_smiles}`\n\n"
-                        f"**{nan_message}**\n\n"
-                        f"**{novelty_status}**\n\n"
-                        f"**Suggested Properties:**\n{suggested_properties_text}",
-                        img,
-                        gif_path,
-                        standardized_smiles,
-                        properties,
-                        suggested_properties
-                    )
-        except Exception as e:
-            print(f"Attempt {attempt + 1} failed: {e}")
-            continue
-
-    # If all attempts fail
-    return (
-        f"**Generation failed:** Could not generate a valid molecule after {repeating_time} attempts.\n\n**{nan_message}**",
-        None,
-        None,
-        "",
-        [],
-        {}
+    # Model Selection
+    model_choice = gr.Radio(
+        choices=list(model_name_mapping.values()),
+        label="Model Zoo",
+        # value="Graph DiT (trained on labeled + unlabeled)"
+        value="Graph DiT (trained on labeled)"
     )
 
-# --------------------------- Feedback Processing --------------------------- #
+    # Model Description Accordion
+    with gr.Accordion("🔍 Model Description", open=False):
+        gr.Markdown("""
+        ### GraphDiT: Graph Diffusion Transformer
 
-def process_feedback(checkbox_value, smiles, properties, suggested_properties):
-    """
-    Process user feedback. If the user finds the polymer interesting,
-    log it accordingly.
-    """
-    if checkbox_value and smiles:
-        save_interesting_log(smiles, properties, suggested_properties)
-        return "Thank you for your feedback! This polymer has been saved to our interesting polymers log."
-    return "Thank you for your feedback!"
-
-# --------------------------- Model Switching --------------------------- #
-
-def switch_model(choice):
-    """Switch the model based on user selection."""
-    internal_name = next(key for key, value in MODEL_NAME_MAPPING.items() if value == choice)
-    return load_model(internal_name)
-
-# --------------------------- Gradio Interface Setup --------------------------- #
-
-def create_gradio_interface():
-    """Create and return the Gradio Blocks interface."""
-    with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
-        # Navigation Bar
-        with gr.Row(elem_id="navbar"):
-            gr.Markdown("""
-            <div style="text-align: center;">
-                <h1>🔗🔬 Polymer Design with GraphDiT</h1>
-                <div style="display: flex; gap: 20px; justify-content: center; align-items: center; margin-top: 10px;">
-                    <a href="https://github.com/liugangcode/Graph-DiT" target="_blank" style="display: flex; align-items: center; gap: 5px; text-decoration: none; color: inherit;">
-                        <img src="https://img.icons8.com/ios-glyphs/30/000000/github.png" alt="GitHub" />
-                        <span>View Code</span>
-                    </a>
-                    <a href="https://arxiv.org/abs/2401.13858" target="_blank" style="text-decoration: none; color: inherit;">
-                        📄 View Paper
-                    </a>
-                </div>
-            </div>
-            """)
+        GraphDiT is a graph diffusion model designed for targeted molecular generation. It employs a conditional diffusion process to iteratively refine molecular structures based on user-specified properties.
 
-        # Main Description
-        gr.Markdown("""
-        ## Introduction
+        We have collected a labeled polymer database for gas permeability from [Membrane Database](https://research.csiro.au/virtualscreening/membrane-database-polymer-gas-separation-membranes/). Additionally, we utilize unlabeled polymer structures from [PolyInfo](https://polymer.nims.go.jp/).
+
+        The gas permeability ranges from 0 to over ten thousand, with only hundreds of labeled data points, making this task particularly challenging.
+
+        We are actively working on improving the model. We welcome any feedback regarding model usage or suggestions for improvement.
 
-        Input the desired gas barrier properties for CH₄, CO₂, H₂, N₂, and O₂ to generate novel polymer structures. The results are visualized as molecular graphs and represented by SMILES strings if they are successfully generated. **Note:** Gas barrier values set to 0 will be treated as `NaN` (unconditionally). If the generation fails, please retry or increase the number of repetition attempts.
+        #### Currently, we have two variants of Graph DiT:
+        - **Graph DiT (trained on labeled + unlabeled)**: This model uses both labeled and unlabeled data for training, potentially leading to more diverse/novel polymer generation.
+        - **Graph DiT (trained on labeled)**: This model is trained exclusively on labeled data, which may result in higher validity but potentially less diverse/novel outputs.
         """)
 
-        # Model Selection
-        model_choice = gr.Radio(
-            choices=list(MODEL_NAME_MAPPING.values()),
-            label="Model Zoo",
-            value=MODEL_NAME_MAPPING["model_labeled"]
-        )
-
-        # Model Description Accordion
-        with gr.Accordion("🔍 Model Description", open=False):
-            gr.Markdown("""
-            ### GraphDiT: Graph Diffusion Transformer
-
-            GraphDiT is a graph diffusion model designed for targeted molecular generation. It employs a conditional diffusion process to iteratively refine molecular structures based on user-specified properties.
-
-            We have collected a labeled polymer database for gas permeability from [Membrane Database](https://research.csiro.au/virtualscreening/membrane-database-polymer-gas-separation-membranes/). Additionally, we utilize unlabeled polymer structures from [PolyInfo](https://polymer.nims.go.jp/).
-
-            The gas permeability ranges from 0 to over ten thousand, with only hundreds of labeled data points, making this task particularly challenging.
-
-            We are actively working on improving the model. We welcome any feedback regarding model usage or suggestions for improvement.
-
-            #### Currently, we have two variants of Graph DiT:
-            - **Graph DiT (trained on labeled + unlabeled)**: This model uses both labeled and unlabeled data for training, potentially leading to more diverse/novel polymer generation.
-            - **Graph DiT (trained on labeled)**: This model is trained exclusively on labeled data, which may result in higher validity but potentially less diverse/novel outputs.
-            """)
-
-        # Citation Accordion
-        with gr.Accordion("📄 Citation", open=False):
-            gr.Markdown("""
-            If you use this model or interface useful, please cite the following paper:
-            ```bibtex
-            @article{graphdit2024,
-              title={Graph Diffusion Transformers for Multi-Conditional Molecular Generation},
-              author={Liu, Gang and Xu, Jiaxin and Luo, Tengfei and Jiang, Meng},
-              journal={NeurIPS},
-              year={2024},
-            }
-            ```
-            """)
-
-        # Initialize Model State
-        model_state = gr.State(load_model("model_labeled"))
-
-        # Property Inputs
-        with gr.Row():
-            CH4_input = gr.Slider(
-                minimum=0,
-                maximum=property_ranges['CH4'][1],
-                value=2.5,
-                label=f"CH₄ (Barrier) [0-{property_ranges['CH4'][1]:.1f}]"
-            )
-            CO2_input = gr.Slider(
-                minimum=0,
-                maximum=property_ranges['CO2'][1],
-                value=15.4,
-                label=f"CO₂ (Barrier) [0-{property_ranges['CO2'][1]:.1f}]"
-            )
-            H2_input = gr.Slider(
-                minimum=0,
-                maximum=property_ranges['H2'][1],
-                value=21.0,
-                label=f"H₂ (Barrier) [0-{property_ranges['H2'][1]:.1f}]"
-            )
-            N2_input = gr.Slider(
-                minimum=0,
-                maximum=property_ranges['N2'][1],
-                value=1.5,
-                label=f"N₂ (Barrier) [0-{property_ranges['N2'][1]:.1f}]"
-            )
-            O2_input = gr.Slider(
-                minimum=0,
-                maximum=property_ranges['O2'][1],
-                value=2.8,
-                label=f"O��� (Barrier) [0-{property_ranges['O2'][1]:.1f}]"
-            )
+    # Citation Accordion
+    with gr.Accordion("📄 Citation", open=False):
+        gr.Markdown("""
+        If you use this model or interface useful, please cite the following paper:
+        ```bibtex
+        @article{graphdit2024,
+          title={Graph Diffusion Transformers for Multi-Conditional Molecular Generation},
+          author={Liu, Gang and Xu, Jiaxin and Luo, Tengfei and Jiang, Meng},
+          journal={NeurIPS},
+          year={2024},
+        }
+        ```
+        """)
 
-        # Generation Parameters
-        with gr.Row():
-            guidance_scale = gr.Slider(
-                minimum=1,
-                maximum=3,
-                value=2,
-                label="Guidance Scale from Properties"
-            )
-            num_nodes = gr.Slider(
-                minimum=0,
-                maximum=50,
-                step=1,
-                value=0,
-                label="Number of Nodes (0 for Random, Larger Graphs Take More Time)"
-            )
-            repeating_time = gr.Slider(
-                minimum=1,
-                maximum=10,
-                step=1,
-                value=3,
-                label="Repetition Until Success"
-            )
-            num_chain_steps = gr.Slider(
-                minimum=0,
-                maximum=499,
-                step=1,
-                value=50,
-                label="Number of Diffusion Steps to Visualize (Larger Numbers Take More Time)"
-            )
-            fps = gr.Slider(
-                minimum=0.25,
-                maximum=10,
-                step=0.25,
-                value=5,
-                label="Frames Per Second"
-            )
+    model_state = gr.State(lambda: load_model("model_labeled"))
+
+    with gr.Row():
+        CH4_input = gr.Slider(0, property_ranges['CH4'][1], value=2.5, label=f"CH₄ (Barrier) [0-{property_ranges['CH4'][1]:.1f}]")
+        CO2_input = gr.Slider(0, property_ranges['CO2'][1], value=15.4, label=f"CO₂ (Barrier) [0-{property_ranges['CO2'][1]:.1f}]")
+        H2_input = gr.Slider(0, property_ranges['H2'][1], value=21.0, label=f"H₂ (Barrier) [0-{property_ranges['H2'][1]:.1f}]")
+        N2_input = gr.Slider(0, property_ranges['N2'][1], value=1.5, label=f"N₂ (Barrier) [0-{property_ranges['N2'][1]:.1f}]")
+        O2_input = gr.Slider(0, property_ranges['O2'][1], value=2.8, label=f"O₂ (Barrier) [0-{property_ranges['O2'][1]:.1f}]")
+
+    with gr.Row():
+        guidance_scale = gr.Slider(1, 3, value=2, label="Guidance Scale from Properties")
+        num_nodes = gr.Slider(0, 50, step=1, value=0, label="Number of Nodes (0 for Random, Larger Graphs Take More Time)")
+        repeating_time = gr.Slider(1, 10, step=1, value=3, label="Repetition Until Success")
+        num_chain_steps = gr.Slider(0, 499, step=1, value=50, label="Number of Diffusion Steps to Visualize (Larger Numbers Take More Time)")
+        fps = gr.Slider(0.25, 10, step=0.25, value=5, label="Frames Per Second")
+
+    with gr.Row():
+        random_btn = gr.Button("🔀 Randomize Properties (from Labeled Data)")
+        generate_btn = gr.Button("🚀 Generate Polymer")
+
+    with gr.Row():
+        result_text = gr.Textbox(label="📝 Generation Result")
+        result_image = gr.Image(label="Final Molecule Visualization", type="pil")
+        result_gif = gr.Image(label="Generation Process Visualization", type="filepath", format="gif")
+
+    with gr.Row() as feedback_row:
+        feedback_btn = gr.Button("🌟 I think this polymer is interesting!", visible=True, interactive=False)
+        feedback_result = gr.Textbox(label="Feedback Result", visible=False)
+    
+    # Add model switching functionality
+    def switch_model(choice):
+        # Convert display name back to internal name
+        internal_name = next(key for key, value in model_name_mapping.items() if value == choice)
+        return load_model(internal_name)
+
+    model_choice.change(switch_model, inputs=[model_choice], outputs=[model_state])
+
+    # Hidden components to store generation data
+    hidden_smiles = gr.Textbox(visible=False)
+    hidden_properties = gr.JSON(visible=False)
+    hidden_suggested_properties = gr.JSON(visible=False)
+    
+    # Set up event handlers
+    random_btn.click(
+        set_random_properties,
+        outputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input]
+    )
 
-        # Action Buttons
-        with gr.Row():
-            random_btn = gr.Button("🔀 Randomize Properties (from Labeled Data)")
-            generate_btn = gr.Button("🚀 Generate Polymer")
-
-        # Results Display
-        with gr.Row():
-            result_text = gr.Textbox(label="📝 Generation Result", lines=10)
-            result_image = gr.Image(label="Final Molecule Visualization", type="pil")
-            result_gif = gr.Image(label="Generation Process Visualization", type="filepath", format="gif")
-
-        # Feedback Section
-        with gr.Row():
-            feedback_btn = gr.Button("🌟 I think this polymer is interesting!", interactive=False)
-            feedback_result = gr.Textbox(label="Feedback Result", visible=False)
-
-        # Hidden Components to Store Generation Data
-        hidden_smiles = gr.Textbox(visible=False)
-        hidden_properties = gr.JSON(visible=False)
-        hidden_suggested_properties = gr.JSON(visible=False)
-
-        # Event Handlers
-
-        # Model Selection Change
-        model_choice.change(
-            switch_model,
-            inputs=[model_choice],
-            outputs=[model_state]
-        )
-
-        # Randomize Properties Button
-        random_btn.click(
-            random_properties,
-            outputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input]
-        )
-
-        # Generate Polymer Button
-        generate_btn.click(
-            generate_graph,
-            inputs=[
-                CH4_input, CO2_input, H2_input, N2_input, O2_input,
-                guidance_scale, num_nodes, repeating_time,
-                model_state, num_chain_steps, fps
-            ],
-            outputs=[
-                result_text, result_image, result_gif,
-                hidden_smiles, hidden_properties, hidden_suggested_properties
-            ]
-        ).then(
-            lambda text, img, gif, smiles, props, sugg_props: (
-                smiles if text.startswith("**Generated polymer SMILES:**") else "",
-                json.dumps(numpy_to_python(props)),
-                json.dumps(numpy_to_python(sugg_props)),
-                gr.Button(interactive=text.startswith("**Generated polymer SMILES:**"))
-            ),
-            inputs=[
-                result_text, result_image, result_gif,
-                hidden_smiles, hidden_properties, hidden_suggested_properties
-            ],
-            outputs=[hidden_smiles, hidden_properties, hidden_suggested_properties, feedback_btn]
-        )
-
-        # Feedback Button Click
-        feedback_btn.click(
-            process_feedback,
-            inputs=[gr.Checkbox(label="Interested?", value=True, visible=False), hidden_smiles, hidden_properties, hidden_suggested_properties],
-            outputs=[feedback_result]
-        ).then(
-            lambda: gr.Button(interactive=False),
-            outputs=[feedback_btn]
-        )
-
-        # # Define the reset_feedback function
-        # def reset_feedback():
-        #     return gr.Button(interactive=False)
-
-        # CH4_input.change(reset_feedback, outputs=[feedback_btn])
-        # CO2_input.change(reset_feedback, outputs=[feedback_btn])
-        # H2_input.change(reset_feedback, outputs=[feedback_btn])
-        # N2_input.change(reset_feedback, outputs=[feedback_btn])
-        # O2_input.change(reset_feedback, outputs=[feedback_btn])
-        # random_btn.click(reset_feedback, outputs=[feedback_btn])
-
-    return iface
-
-# --------------------------- Main Execution --------------------------- #
+    generate_btn.click(
+        on_generate,
+        inputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps],
+        outputs=[result_text, result_image, result_gif, hidden_smiles, hidden_properties, hidden_suggested_properties, feedback_btn]
+    )
 
+    feedback_btn.click(
+        process_feedback,
+        inputs=[gr.Checkbox(value=True, visible=False), hidden_smiles, hidden_properties, hidden_suggested_properties],
+        outputs=[feedback_result]
+    ).then(
+        lambda: gr.Button(interactive=False),
+        outputs=[feedback_btn]
+    )
+    
+    CH4_input.change(reset_feedback_button, outputs=[feedback_btn])
+    CO2_input.change(reset_feedback_button, outputs=[feedback_btn])
+    H2_input.change(reset_feedback_button, outputs=[feedback_btn])
+    N2_input.change(reset_feedback_button, outputs=[feedback_btn])
+    O2_input.change(reset_feedback_button, outputs=[feedback_btn])
+    random_btn.click(reset_feedback_button, outputs=[feedback_btn])
+
+# Launch the interface
 if __name__ == "__main__":
-    interface = create_gradio_interface()
-    interface.launch(share=False)
+    # iface.launch(share=True)
+    iface.launch(share=False)