update

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from huggingface_hub import HfApi, get_collection, list_collections
-from utils import MolecularPropertyPredictionModel, task_types, dataset_descriptions
+from utils import MolecularPropertyPredictionModel, dataset_task_types, dataset_descriptions, dataset_property_names, dataset_property_names_to_dataset
 import pandas as pd
 import os
 
@@ -12,22 +12,26 @@ def get_models():
         if item.item_type == "model":
             item_name = item.item_id.split("/")[-1]
             models[item_name] = item.item_id
-            assert item_name in task_types, f"{item_name} is not in the task_types"
+            assert item_name in dataset_task_types, f"{item_name} is not in the task_types"
             assert item_name in dataset_descriptions, f"{item_name} is not in the dataset_descriptions"
     
     return models
 
 candidate_models = get_models()
-properties = list(candidate_models.keys())
+properties = [dataset_property_names[item] for item in candidate_models.keys()]
+property_names = list(candidate_models.keys())
 model = MolecularPropertyPredictionModel(candidate_models)
 
 def get_description(property_name):
-    return dataset_descriptions[property_name]
+    property_id = dataset_property_names_to_dataset[property_name]
+    return dataset_descriptions[property_id]
 
 def predict_single_label(smiles, property_name):
+    property_id = dataset_property_names_to_dataset[property_name]
+
     try:
-        adapter_id = candidate_models[property_name]
-        info = model.swith_adapter(property_name, adapter_id)
+        adapter_id = candidate_models[property_id]
+        info = model.swith_adapter(property_id, adapter_id)
 
         running_status = None
         if info == "keep":
@@ -45,7 +49,7 @@ def predict_single_label(smiles, property_name):
             return "NA", running_status
     
         #prediction = model.predict(smiles, property_name, adapter_id)
-        prediction = model.predict_single_smiles(smiles, task_types[property_name])
+        prediction = model.predict_single_smiles(smiles, dataset_task_types[property_id])
         if prediction is None:
             return "NA", "Invalid SMILES string"
     
@@ -60,9 +64,10 @@ def predict_single_label(smiles, property_name):
     return prediction, "Prediction is done"
 
 def predict_file(file, property_name):
+    property_id = dataset_property_names_to_dataset[property_name]
     try:
-        adapter_id = candidate_models[property_name]
-        info = model.swith_adapter(property_name, adapter_id)
+        adapter_id = candidate_models[property_id]
+        info = model.swith_adapter(property_id, adapter_id)
 
         running_status = None
         if info == "keep":
@@ -81,7 +86,7 @@ def predict_file(file, property_name):
     
         df = pd.read_csv(file)
         # we have already checked the file contains the "smiles" column
-        df = model.predict_file(df, task_types[property_name])
+        df = model.predict_file(df, dataset_task_types[property_id])
         # we should save this file to the disk to be downloaded
         # rename the file to have "_prediction" suffix
         prediction_file = file.replace(".csv", "_prediction.csv") if file.endswith(".csv") else file.replace(".smi", "_prediction.csv")
@@ -157,10 +162,12 @@ def build_inference():
     with gr.Blocks() as demo:
         # first row - Dropdown input
         #with gr.Row():
-        dropdown = gr.Dropdown(properties, label="Property", value=properties[0])
+        print(property_names[0].lower())
+        print(properties)
+        dropdown = gr.Dropdown(properties, label="Property", value=dataset_property_names[property_names[0].lower()])
         description_box = gr.Textbox(label="Property description", lines=5,
                                      interactive=False,
-                                     value=dataset_descriptions[properties[0]])
+                                     value=dataset_descriptions[property_names[0].lower()])
         # third row - Textbox input and prediction label
         with gr.Row(equal_height=True):
             with gr.Column():