Upload app.py with huggingface_hub

app.py

@@ -1,142 +1,347 @@
-# import spaces
+import spaces
 import gradio as gr
+
 import torch
-import torch.nn as nn
+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
+
 from rdkit import Chem
 from rdkit.Chem import Draw
 
-from graph_decoder.diffusion_model import GraphDiT
+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')
 
-ATOM_SYMBOLS = ['C', 'N', 'O', 'H']
+all_properties = ['CH4', 'CO2', 'H2', 'N2', 'O2']
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Initialize evaluators
+evaluators = {prop: Evaluator(f'evaluators/{prop}.joblib', prop) for prop in all_properties}
 
-path = 'model_labeled'
-model = GraphDiT(
-        model_config_path=f"{path}/config.yaml",
-        data_info_path=f"{path}/data.meta.json",
-        model_dtype=torch.float32
-        )
-model.to(device)
+# 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 generate_random_smiles(length=10):
-    return ''.join(random.choices(ATOM_SYMBOLS, k=length))
+# Create a temporary directory for GIFs
+temp_dir = tempfile.mkdtemp(prefix="polymer_gifs_")
 
-# @spaces.GPU
-def generate_polymer(CH4, CO2, H2, N2, O2, guidance_scale):
-    properties = torch.tensor([CH4, CO2, H2, N2, O2], dtype=torch.float32).unsqueeze(0)
-    
-    print('in generate_polymer')
+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:
-        # Generate a random SMILES string (this is a placeholder)
-        generated_molecule = generate_random_smiles()
-        model.generate(properties, device)
-        
-        mol = Chem.MolFromSmiles(generated_molecule)
-        if mol is not None:
-            standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
-            img = Draw.MolToImage(mol)
-            return standardized_smiles, img
+        os.rmdir(temp_dir)
     except Exception as e:
-        print(f"Error in generation: {e}")
+        print(f"Error deleting temporary directory {temp_dir}: {e}")
+
+# Register the cleanup function to be called on exit
+atexit.register(cleanup_temp_files)
+
+def random_properties():
+    return known_labels[all_properties].sample(1).values.tolist()[0]
+
+def load_model(model_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)
+
+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")
+    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):
+    model, device = model_state
+    
+    properties = [CH4, CO2, H2, N2, O2]
     
-    return "Generation failed", None
+    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"
 
-# Create the Gradio interface
-with gr.Blocks(title="Simplified Polymer Design") as iface:
-    gr.Markdown("## Polymer Design with Random Neural Network")
+    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 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):
+    if isinstance(obj, np.integer):
+        return int(obj)
+    elif isinstance(obj, np.floating):
+        return float(obj)
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    elif isinstance(obj, list):
+        return [numpy_to_python(item) for item in obj]
+    elif isinstance(obj, dict):
+        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):
+    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)
+
+# ADD THIS FUNCTION
+def reset_feedback_button():
+    return gr.Button(interactive=False)
+
+# 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>
+        """)
+
+    # Main Description
+    gr.Markdown("""
+    ## Introduction
+
+    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.
+    """)
+
+    # 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)"
+    )
+
+    # 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},
+        }
+        ```
+        """)
+
+    model_state = gr.State(lambda: load_model("model_labeled"))
+
     with gr.Row():
-        CH4_input = gr.Slider(0, 100, value=2.5, label="CH₄ (Barrier)")
-        CO2_input = gr.Slider(0, 100, value=15.4, label="CO₂ (Barrier)")
-        H2_input = gr.Slider(0, 100, value=21.0, label="H₂ (Barrier)")
-        N2_input = gr.Slider(0, 100, value=1.5, label="N₂ (Barrier)")
-        O2_input = gr.Slider(0, 100, value=2.8, label="O₂ (Barrier)")
-        guidance_scale = gr.Slider(1, 3, value=2, label="Guidance Scale")
+        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}]")
 
-    generate_btn = gr.Button("Generate Polymer")
+    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():
-        result_smiles = gr.Textbox(label="Generated SMILES")
-        result_image = gr.Image(label="Molecule Visualization", type="pil")
+        random_btn = gr.Button("🔀 Randomize Properties (from Labeled Data)")
+        generate_btn = gr.Button("🚀 Generate Polymer")
 
-    generate_btn.click(
-        generate_polymer,
-        inputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input, guidance_scale],
-        outputs=[result_smiles, result_image]
-    )
+    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")
 
-if __name__ == "__main__":
-    iface.launch()
-
-# import spaces
-# import gradio as gr
-# import torch
-# from rdkit import Chem
-# from rdkit.Chem import Draw
-# # from graph_decoder.diffusion_model import GraphDiT
-
-# # Load the model
-# def load_graph_decoder(path='model_labeled'):
-#     model = GraphDiT(
-#         model_config_path=f"{path}/config.yaml",
-#         data_info_path=f"{path}/data.meta.json",
-#         model_dtype=torch.float32,
-#     )
-#     model.init_model(path)
-#     model.disable_grads()
-#     return model
-
-# # model = load_graph_decoder()
-# # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# @spaces.GPU
-# def generate_polymer(CH4, CO2, H2, N2, O2, guidance_scale):
-#     properties = [CH4, CO2, H2, N2, O2]
+    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)
     
-#     try:
-#         model = load_graph_decoder()
-#         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#         model.to(device)
-#         print('enter function')
-#         generated_molecule, _ = model.generate(properties, device=device, guide_scale=guidance_scale)
-        
-#         if generated_molecule is not None:
-#             mol = Chem.MolFromSmiles(generated_molecule)
-#             if mol is not None:
-#                 standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
-#                 img = Draw.MolToImage(mol)
-#                 return standardized_smiles, img
-#     except Exception as e:
-#         print(f"Error in generation: {e}")
+    # 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)
     
-#     return "Generation failed", None
+    # Set up event handlers
+    random_btn.click(
+        set_random_properties,
+        outputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input]
+    )
 
-# # Create the Gradio interface
-# with gr.Blocks(title="Simplified Polymer Design") as iface:
-#     gr.Markdown("## Polymer Design with GraphDiT")
+    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]
+    )
     
-#     with gr.Row():
-#         CH4_input = gr.Slider(0, 100, value=2.5, label="CH₄ (Barrier)")
-#         CO2_input = gr.Slider(0, 100, value=15.4, label="CO₂ (Barrier)")
-#         H2_input = gr.Slider(0, 100, value=21.0, label="H₂ (Barrier)")
-#         N2_input = gr.Slider(0, 100, value=1.5, label="N₂ (Barrier)")
-#         O2_input = gr.Slider(0, 100, value=2.8, label="O₂ (Barrier)")
-#         guidance_scale = gr.Slider(1, 3, value=2, label="Guidance Scale")
-
-#     generate_btn = gr.Button("Generate Polymer")
-
-#     with gr.Row():
-#         result_smiles = gr.Textbox(label="Generated SMILES")
-#         result_image = gr.Image(label="Molecule Visualization", type="pil")
-
-#     generate_btn.click(
-#         generate_polymer,
-#         inputs=[CH4_input, CO2_input, H2_input, N2_input, O2_input, guidance_scale],
-#         outputs=[result_smiles, result_image]
-#     )
-
-# if __name__ == "__main__":
-#     iface.launch()
+    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)