Upload app.py with huggingface_hub

app.py

@@ -1,168 +1,137 @@
 import spaces
-import gradio as gr
-import torch
 
-import numpy as np
-import pandas as pd
-import random
-import io
-import imageio
+# Standard Libraries
 import os
+import io
+import csv
+import json
+import glob
+import random
 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
 
-# Load the CSV data
-known_labels = pd.read_csv('data/known_labels.csv')
-knwon_smiles = pd.read_csv('data/known_polymers.csv')
+# --------------------------- 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()
+
+# --------------------------- 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
 
-all_properties = ['CH4', 'CO2', 'H2', 'N2', 'O2']
+evaluators = initialize_evaluators(ALL_PROPERTIES, EVALUATORS_DIR)
 
-# Initialize evaluators
-evaluators = {prop: Evaluator(f'evaluators/{prop}.joblib', prop) for prop in all_properties}
+# --------------------------- Property Ranges --------------------------- #
 
-# Get min and max values for each property
-property_ranges = {prop: (known_labels[prop].min(), known_labels[prop].max()) for prop in all_properties}
+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}
 
-# Create a temporary directory for GIFs
-temp_dir = tempfile.mkdtemp(prefix="polymer_gifs_")
+property_ranges = get_property_ranges(known_labels, ALL_PROPERTIES)
+
+# --------------------------- Temporary Directory Setup --------------------------- #
+
+temp_dir = tempfile.mkdtemp(prefix=GIF_TEMP_PREFIX)
 
 def cleanup_temp_files():
     """Clean up temporary GIF files on exit."""
-    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:
+        for file in glob.glob(os.path.join(temp_dir, "*.gif")):
+            os.remove(file)
         os.rmdir(temp_dir)
     except Exception as e:
-        print(f"Error deleting temporary directory {temp_dir}: {e}")
+        print(f"Error during cleanup: {e}")
 
-# Register the cleanup function to be called on exit
 atexit.register(cleanup_temp_files)
 
+# --------------------------- Utility Functions --------------------------- #
+
 def random_properties():
-    return known_labels[all_properties].sample(1).values.tolist()[0]
+    """Select a random set of properties from known labels."""
+    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)
-
-# Create a flagged folder if it doesn't exist
-flagged_folder = "flagged"
-os.makedirs(flagged_folder, exist_ok=True)
+    model.to(device)
+    return model, device
 
 def save_interesting_log(smiles, properties, suggested_properties):
-    """Save interesting polymer data to a CSV file."""
-    log_file = os.path.join(flagged_folder, "log.csv")
+    """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)
     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)
-
-@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"
 
-    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
+    fieldnames = ['timestamp', 'smiles'] + ALL_PROPERTIES + [f'suggested_{prop}' for prop in ALL_PROPERTIES]
 
-def set_random_properties():
-    return random_properties()
+    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}")
 
-# 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 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 numpy_to_python(obj):
+    """Convert NumPy objects to native Python types."""
     if isinstance(obj, np.integer):
         return int(obj)
     elif isinstance(obj, np.floating):
@@ -175,175 +144,326 @@ 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)]
 
-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)
+# --------------------------- 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]
 
-# ADD THIS FUNCTION
-def reset_feedback_button():
-    return gr.Button(interactive=False)
+    # Handle NaN-like values
+    properties = [None if is_nan_like(prop) else prop for prop in properties]
 
-# 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>
-        """)
+    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")
 
-    # Main Description
-    gr.Markdown("""
-    ## Introduction
+    num_nodes = None if num_nodes == 0 else num_nodes
 
-    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.
-    """)
+    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
+            )
 
-    # 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)"
+            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 Description Accordion
-    with gr.Accordion("🔍 Model Description", open=False):
-        gr.Markdown("""
-        ### GraphDiT: Graph Diffusion Transformer
+# --------------------------- Feedback Processing --------------------------- #
 
-        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.
+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>
+            """)
 
-        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/).
+        # Main Description
+        gr.Markdown("""
+        ## Introduction
 
-        The gas permeability ranges from 0 to over ten thousand, with only hundreds of labeled data points, making this task particularly challenging.
+        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.
+        """)
 
-        We are actively working on improving the model. We welcome any feedback regarding model usage or suggestions for improvement.
+        # 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}]"
+            )
 
-        #### 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.
-        """)
+        # 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"
+            )
 
-    # 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},
-        }
-        ```
-        """)
+        # 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]
+        )
+
+        # Reset Feedback Button on Input Changes
+        for input_component in [CH4_input, CO2_input, H2_input, N2_input, O2_input, random_btn]:
+            input_component.change(
+                lambda: None,
+                outputs=[feedback_btn],
+                _js="() => feedback_btn.interactive = false"
+            )
 
-    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]
-    )
+    return iface
 
-    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]
-    )
+# --------------------------- Main Execution --------------------------- #
 
-    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__":
-    # iface.launch(share=True)
-    iface.launch(share=False)
+    interface = create_gradio_interface()
+    interface.launch(share=False)