Change globals to state. Change transformers version. Delete code history comments.

Dockerfile

@@ -3,6 +3,11 @@ FROM python:3.9.7
 WORKDIR /app
 COPY requirements.txt .
 RUN pip install -r requirements.txt
+# preload models
+RUN python -c '\
+from transformers import BartForConditionalGeneration, AutoTokenizer;\
+AutoTokenizer.from_pretrained("ibm/materials.selfies-ted");\
+BartForConditionalGeneration.from_pretrained("ibm/materials.selfies-ted")'
 COPY . .
 
 CMD ["python", "app.py"]

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
 import matplotlib.pyplot as plt
 from PIL import Image
 from rdkit.Chem import Descriptors, QED, Draw
@@ -25,23 +24,6 @@ import os
 
 os.environ["OMP_MAX_ACTIVE_LEVELS"] = "1"
 
-# my_theme = gr.Theme.from_hub("ysharma/steampunk")
-# my_theme = gr.themes.Glass()
-
-"""
-# カスタムテーマ設定
-theme = gr.themes.Default().set(
-    body_background_fill="#000000",  # 背景色を黒に設定
-    text_color="#FFFFFF",            # テキスト色を白に設定
-)
-"""
-"""
-import sys
-sys.path.append("models")
-sys.path.append("../models")
-sys.path.append("../")"""
-
-
 # Get the current file's directory
 base_dir = os.path.dirname(__file__)
 print("Base Dir : ", base_dir)
@@ -139,7 +121,6 @@ def load_image(path):
         pass
 
 
-
 # Function to handle image selection
 def handle_image_selection(image_key):
     smiles = smiles_image_mapping[image_key]["smiles"]
@@ -171,9 +152,6 @@ def calculate_tanimoto(smiles1, smiles2):
     return None
 
 
-#with open("models/selfies_model/bart-2908.pickle", "rb") as input_file:
-#    gen_model, gen_tokenizer = pickle.load(input_file)
-
 gen_tokenizer = AutoTokenizer.from_pretrained("ibm/materials.selfies-ted")
 gen_model = BartForConditionalGeneration.from_pretrained("ibm/materials.selfies-ted")
 
@@ -201,11 +179,6 @@ def encode(selfies):
     attention_mask = encoding['attention_mask']
     outputs = gen_model.model.encoder(input_ids=input_ids, attention_mask=attention_mask)
     model_output = outputs.last_hidden_state
-
-    """input_mask_expanded = attention_mask.unsqueeze(-1).expand(model_output.size()).float()
-    sum_embeddings = torch.sum(model_output * input_mask_expanded, 1)
-    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-    model_output = sum_embeddings / sum_mask"""
     return model_output, attention_mask
 
 
@@ -226,16 +199,6 @@ def generate_canonical(smiles):
             if gen_mol != Chem.MolToSmiles(Chem.MolFromSmiles(smiles)): break
         else:
             print('Abnormal molecule:', gen[0])
-            gen_mols = []
-            for sel in gen[0].split('.'):
-                mol = Chem.MolFromSmiles(sel)
-                if mol:
-                    mol = Chem.MolToSmiles(mol)
-                    if mol != Chem.MolToSmiles(Chem.MolFromSmiles(smiles)):
-                        gen_mols.append(mol)
-            if len(gen_mols) > 0:
-                gen_mol = '.'.join(gen_mols)
-                break
 
     if gen_mol:
         # Calculate properties for ref and gen molecules
@@ -262,7 +225,7 @@ def generate_canonical(smiles):
 
 
 # Function to display evaluation score
-def display_eval(selected_models, dataset, task_type, downstream, fusion_type):
+def display_eval(selected_models, dataset, task_type, downstream, fusion_type, state):
     result = None
 
     try:
@@ -278,68 +241,47 @@ def display_eval(selected_models, dataset, task_type, downstream, fusion_type):
         params = None
 
 
-
-
     try:
         if not selected_models:
-            return "Please select at least one enabled model."
-
-        if task_type == "Classification":
-            global roc_auc, fpr, tpr, x_batch, y_batch
-        elif task_type == "Regression":
-            global RMSE, y_batch_test, y_prob
+            return "Please select at least one enabled model.", state
 
         if len(selected_models) > 1:
             if task_type == "Classification":
-                #result, roc_auc, fpr, tpr, x_batch, y_batch = fm4m.multi_modal(model_list=selected_models,
-                #                                                               downstream_model="XGBClassifier",
-                #                                                               dataset=dataset.lower())
                 if downstream_model == "Default Settings":
                     downstream_model = "DefaultClassifier"
                     params = None
-                result, roc_auc, fpr, tpr, x_batch, y_batch = fm4m.multi_modal(model_list=selected_models,
+                result, state["roc_auc"], state["fpr"], state["tpr"], state["x_batch"], state["y_batch"] = fm4m.multi_modal(model_list=selected_models,
                                                                                                downstream_model=downstream_model,
                                                                                                params = params,
                                                                                                dataset=dataset)
 
             elif task_type == "Regression":
-                #result, RMSE, y_batch_test, y_prob = fm4m.multi_modal(model_list=selected_models,
-                #                                                      downstream_model="XGBRegressor",
-                #                                                      dataset=dataset.lower())
-
                 if downstream_model == "Default Settings":
                     downstream_model = "DefaultRegressor"
                     params = None
 
-                result, RMSE, y_batch_test, y_prob, x_batch, y_batch = fm4m.multi_modal(model_list=selected_models,
+                result, state["RMSE"], state["y_batch_test"], state["y_prob"], state["x_batch"], state["y_batch"] = fm4m.multi_modal(model_list=selected_models,
                                                                       downstream_model=downstream_model,
                                                                       params=params,
                                                                       dataset=dataset)
 
         else:
             if task_type == "Classification":
-                #result, roc_auc, fpr, tpr, x_batch, y_batch = fm4m.single_modal(model=selected_models[0],
-                #                                                                downstream_model="XGBClassifier",
-                #                                                                dataset=dataset.lower())
                 if downstream_model == "Default Settings":
                     downstream_model = "DefaultClassifier"
                     params = None
 
-                result, roc_auc, fpr, tpr, x_batch, y_batch = fm4m.single_modal(model=selected_models[0],
+                result, state["roc_auc"], state["fpr"], state["tpr"], state["x_batch"], state["y_batch"] = fm4m.single_modal(model=selected_models[0],
                                                                                 downstream_model=downstream_model,
                                                                                 params=params,
                                                                                 dataset=dataset)
 
             elif task_type == "Regression":
-                #result, RMSE, y_batch_test, y_prob = fm4m.single_modal(model=selected_models[0],
-                #                                                       downstream_model="XGBRegressor",
-                #                                                       dataset=dataset.lower())
-
                 if downstream_model == "Default Settings":
                     downstream_model = "DefaultRegressor"
                     params = None
 
-                result, RMSE, y_batch_test, y_prob, x_batch, y_batch = fm4m.single_modal(model=selected_models[0],
+                result, state["RMSE"], state["y_batch_test"], state["y_prob"], state["x_batch"], state["y_batch"] = fm4m.single_modal(model=selected_models[0],
                                                                        downstream_model=downstream_model,
                                                                        params=params,
                                                                        dataset=dataset)
@@ -347,28 +289,20 @@ def display_eval(selected_models, dataset, task_type, downstream, fusion_type):
         if result == None:
             result = "Data & Model Setting is incorrect"
     except Exception as e:
-        return f"An error occurred: {e}"
-    return f"{result}"
+        return f"An error occurred: {e}", state
+    return f"{result}", state
 
 
 # Function to handle plot display
-def display_plot(plot_type):
+def display_plot(plot_type, state):
     fig, ax = plt.subplots()
 
     if plot_type == "Latent Space":
-        global x_batch, y_batch
+        x_batch, y_batch = state.get("x_batch"), state.get("y_batch")
         ax.set_title("T-SNE Plot")
-        # reducer = umap.UMAP(metric='euclidean', n_neighbors=  10, n_components=2, low_memory=True, min_dist=0.1, verbose=False)
-        # features_umap = reducer.fit_transform(x_batch[:500])
-        # x = y_batch.values[:500]
-        # index_0 = [index for index in range(len(x)) if x[index] == 0]
-        # index_1 = [index for index in range(len(x)) if x[index] == 1]
-        class_0 = x_batch  # features_umap[index_0]
-        class_1 = y_batch  # features_umap[index_1]
-
-        """with open("latent_multi_bace.pkl", "rb") as f:
-            class_0, class_1 = pickle.load(f)
-        """
+        class_0 = x_batch
+        class_1 = y_batch
+
         plt.scatter(class_1[:, 0], class_1[:, 1], c='red', label='Class 1')
         plt.scatter(class_0[:, 0], class_0[:, 1], c='blue', label='Class 0')
 
@@ -377,7 +311,7 @@ def display_plot(plot_type):
         ax.set_title('Dataset Distribution')
 
     elif plot_type == "ROC-AUC":
-        global roc_auc, fpr, tpr
+        roc_auc, fpr, tpr = state.get("roc_auc"), state.get("fpr"), state.get("tpr")
         ax.set_title("ROC-AUC Curve")
         try:
             ax.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.4f})')
@@ -392,7 +326,7 @@ def display_plot(plot_type):
         ax.legend(loc='lower right')
 
     elif plot_type == "Parity Plot":
-        global RMSE, y_batch_test, y_prob
+        RMSE, y_batch_test, y_prob = state.get("RMSE"), state.get("y_batch_test"), state.get("y_prob")
         ax.set_title("Parity plot")
 
         # change format
@@ -415,9 +349,6 @@ def display_plot(plot_type):
             print(y_prob)
 
 
-
-
-
         ax.set_xlabel('Actual Values')
         ax.set_ylabel('Predicted Values')
 
@@ -493,6 +424,7 @@ def create_model(model_name, max_depth=None, n_estimators=None, alpha=None, degr
         return "Model not supported."
 
     return f"{model_name} * {model.get_params()}"
+
 def model_selector(model_name):
     # Dynamically return the appropriate hyperparameter components based on the selected model
     if model_name == "XGBClassifier":
@@ -518,10 +450,10 @@ def model_selector(model_name):
         return ()
 
 
-
 # Define the Gradio layout
 # with gr.Blocks(theme=my_theme) as demo:
 with gr.Blocks() as demo:
+    state = gr.State({})
     with gr.Row():
         # Left Column
         with gr.Column():
@@ -530,9 +462,6 @@ with gr.Blocks() as demo:
                 <h3 style="color: #FFFFFF; margin: 0;font-size: 20px;"> Data & Model Setting</h3>
             </div>
             ''')
-            # gr.Markdown("## Data & Model Setting")
-            #dataset_dropdown = gr.Dropdown(choices=datasets, label="Select Dat")
-
             # Dropdown menu for predefined datasets including "Custom Dataset" option
             dataset_selector = gr.Dropdown(label="Select Dataset",
                                            choices=list(predefined_datasets.keys()) + ["Custom Dataset"])
@@ -553,13 +482,10 @@ with gr.Blocks() as demo:
                                                   interactive=False)
 
 
-
                 # Dropdowns for selecting input and output columns for the custom dataset
                 input_column_selector = gr.Dropdown(label="Select Input Column", choices=[], visible=False)
                 output_column_selector = gr.Dropdown(label="Select Output Column", choices=[], visible=False)
 
-                #selected_columns_message = gr.Textbox(label="Selected Columns Info", visible=True)
-
                 # When a dataset is selected, show either file upload fields (for custom) or load predefined datasets
                 dataset_selector.change(handle_dataset_selection,
                                         inputs=dataset_selector,
@@ -594,10 +520,6 @@ with gr.Blocks() as demo:
 
             model_checkbox = gr.CheckboxGroup(choices=models_enabled, label="Select Model")
 
-            # Add disabled checkboxes for GNN and FNN
-            # gnn_checkbox = gr.Checkbox(label="GNN (Disabled)", value=False, interactive=False)
-            # fnn_checkbox = gr.Checkbox(label="FNN (Disabled)", value=False, interactive=False)
-
             task_radiobutton = gr.Radio(choices=["Classification", "Regression"], label="Task Type")
 
             ####### adding hyper parameter tuning ###########
@@ -662,9 +584,7 @@ with gr.Blocks() as demo:
             fusion_radiobutton = gr.Radio(choices=fusion_available, label="Fusion Type")
 
 
-
             eval_button = gr.Button("Train downstream model")
-            #eval_button.style(css_class="custom-button-left")
 
         # Middle Column
         with gr.Column():
@@ -673,23 +593,20 @@ with gr.Blocks() as demo:
                 <h3 style="color: #FFFFFF; margin: 0;font-size: 20px;"> Downstream Task 1: Property Prediction</h3>
             </div>
             ''')
-            # gr.Markdown("## Downstream task Result")
             eval_output = gr.Textbox(label="Train downstream model")
 
             plot_radio = gr.Radio(choices=["ROC-AUC", "Parity Plot", "Latent Space"], label="Select Plot Type")
             plot_output = gr.Plot(label="Visualization")#, height=250, width=250)
 
-            #download_rep = gr.Button("Download representation")
-
             create_log = gr.Button("Store log")
 
             log_table = gr.Dataframe(value=log_df, label="Log of Selections and Results", interactive=False)
 
             eval_button.click(display_eval,
-                              inputs=[model_checkbox, selected_columns_message, task_radiobutton, output, fusion_radiobutton],
-                              outputs=eval_output)
+                              inputs=[model_checkbox, selected_columns_message, task_radiobutton, output, fusion_radiobutton, state],
+                              outputs=[eval_output, state])
 
-            plot_radio.change(display_plot, inputs=plot_radio, outputs=plot_output)
+            plot_radio.change(display_plot, inputs=[plot_radio, state], outputs=plot_output)
 
 
             # Function to gather selected models
@@ -700,9 +617,6 @@ with gr.Blocks() as demo:
 
             create_log.click(evaluate_and_log, inputs=[model_checkbox, dataset_name, task_radiobutton, eval_output],
                              outputs=log_table)
-            #download_rep.click(save_rep, inputs=[model_checkbox, dataset_name, task_radiobutton, eval_output],
-            #                 outputs=None)
-
         # Right Column
         with gr.Column():
             gr.HTML('''
@@ -710,7 +624,6 @@ with gr.Blocks() as demo:
                 <h3 style="color: #FFFFFF; margin: 0;font-size: 20px;"> Downstream Task 2: Molecule Generation</h3>
             </div>
             ''')
-            # gr.Markdown("## Molecular Generation")
             smiles_input = gr.Textbox(label="Input SMILES String")
             image_display = gr.Image(label="Molecule Image", height=250, width=250)
             # Show images for selection
@@ -719,7 +632,6 @@ with gr.Blocks() as demo:
                     choices=list(smiles_image_mapping.keys()),
                     label="Select from sample molecules",
                     value=None,
-                    #item_images=[load_image(smiles_image_mapping[key]["image"]) for key in smiles_image_mapping.keys()]
                 )
                 image_selector.change(load_image, image_selector, image_display)
             generate_button = gr.Button("Generate")
@@ -728,7 +640,6 @@ with gr.Blocks() as demo:
             property_table = gr.Dataframe(label="Molecular Properties Comparison")
 
 
-
             # Handle image selection
             image_selector.change(handle_image_selection, inputs=image_selector, outputs=[smiles_input, image_display])
             smiles_input.change(smiles_to_image, inputs=smiles_input, outputs=image_display)

requirements.txt

@@ -16,7 +16,7 @@ datasets>=2.13.1
 evaluate>=0.4.0
 selfies>=2.1.0
 gradio==3.41.0
-transformers>=4.38
+transformers==4.38.1
 requests>=2.32.2
 urllib3>=2.2.2
 aiohttp>=3.10.2