fix: update input file implementation

Author: mesakhcienet

examples/science/af3-slurm/examples/simple_ipynb_launcher/Ipynb.md

@@ -6,22 +6,29 @@ Begin by uploading your input JSON file, which contains the data required by Alp
 
    <img src="adm/upload_file.png" alt="try to upload file under the alphafold folder" width="400">
 
-When specifying the input file name (e.g., `fold_input.json`) during the configuration of your data pipeline or inference, it's crucial to reference the correct **relative file path**. Your files should be organized like this:
+Ensure your input JSON file is placed in the correct location **relative to your Jupyter Notebook file**. Your project directory should be structured like this:
 
 ```code
 alphafold/
-├── <your_input_file>.json
+├── <uploaded_file>.json
+├── fold_input.json # Example input file (provided)
 └── slurm-rest-api-notebook.ipynb
 ```
 
-The input JSON file should be in the `alphafold` folder.
+When configuring your data pipeline or inference process, use the correct **relative file path to reference the input JSON**. We provide an example input file of data pipeline process named `fold_input.json` to help you get started quickly.
 
-Usage:
+Since both the notebook and the input file are in the same directory (`alphafold/`), you can specify the file like this in your notebook code:
 
 ```python
 input_file= "fold_input.json"
 ```
 
+If you're uploading your own file, simply replace the `input_file` variable with the name of your uploaded file:
+
+```python
+input_file = "<uploaded_file>.json"
+```
+
 **2. Running the Setup Cells (including System Configuration and SLURM Initialization):**
 
 Once the input file is uploaded, execute all the cells in the notebook that are responsible for setting up the environment, installing dependencies, and loading necessary functions. Ensure all these setup cells run without errors before proceeding to the next section.

examples/science/af3-slurm/examples/simple_ipynb_launcher/adm/slurm-rest-api-notebook.ipynb.j2

@@ -205,7 +205,7 @@
     "datapipeline_partition_name = \"{{ default_datapipeline_partition_name }}\"\n",
     "\n",
     "# Assign `input_file` with datapipeline input file path\n",
-    "input_file= \"\" # Replace with your input file path (relative path from the jupyter notebook directory)\n"
+    "datapipeline_input_file= \"\" # Replace with your input file path (relative path from the jupyter notebook directory)\n"
    ]
   },
   {
@@ -217,7 +217,7 @@
     "# Submit a datapipeline job\n",
     "datapipeline_submit_job_response = client.submit_data_pipeline_job(\n",
     "    partition_name=datapipeline_partition_name, \n",
-    "    input_file=input_file\n",
+    "    input_file=datapipeline_input_file\n",
     ")\n",
     "print(\"Submit Job Response:\", json.dumps(datapipeline_submit_job_response,indent=4))"
    ]
@@ -252,10 +252,16 @@
     "# Make sure to use the correct partition name that listed on AF3 System setting section\n",
     "inference_partition_name = \"{{ default_inference_partition_name }}\"\n",
     "\n",
-    "# Assign `input_file` with inference input file path\n",
+    "# Assign 'inference_input_file' with inference input file path.\n",
     "# You can use the output from the datapipeline job as input for inference\n",
     "# or you can use your own input file relative path from the jupyter notebook directory\n",
-    "inference_input_file = datapipeline_submit_job_response[\"datapipeline_output_path\"]\n",
+    "# In this example, we will use the output from the datapipeline job.\n",
+    "# Datapipeline job could returns multiple output files in case the input file is a JSON file with multiple protein sequences\n",
+    "# The output path is a list of dictionaries, each containing the path to the output file\n",
+    "# The first output path is used as input for inference\n",
+    "print(\"[INFO] Extracting output paths from inference job response...\")\n",
+    "print(\"[INFO] Output count:\", len(datapipeline_submit_job_response[\"output_path_list\"]))\n",
+    "inference_input_file = datapipeline_submit_job_response[\"output_path_list\"][0][\"output_path\"]\n",
     "\n",
     "# Assign `enable_unified_memory` with True or False\n",
     "# No changes required to run default settings\n",
@@ -310,6 +316,22 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Extract the information of output paths from the inference job response\n",
+    "# This is the path where the inference results are stored\n",
+    "# You can use this path to read the output files such as CIF and JSON files\n",
+    "# The output path is a list of dictionaries, each containing the path to the output file\n",
+    "# In this example, we will use the first output path\n",
+    "print(\"[INFO] Extracting output paths from inference job response...\")\n",
+    "print(\"[INFO] Output count:\", len(inference_submit_job_response[\"output_path_list\"]))\n",
+    "inference_output_path_info = inference_submit_job_response[\"output_path_list\"][0]"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -318,8 +340,8 @@
    "source": [
     "# Load CIF file\n",
     "# Note: Adjust the file path to your actual CIF file\n",
-    "cif_file_path = inference_submit_job_response[\"cif_output_path\"]\n",
-    "structure, cif_content = read_cif_file(cif_file_path)\n",
+    "model_cif_output_path = inference_output_path_info[\"model_cif\"]\n",
+    "structure, cif_content = read_cif_file(model_cif_output_path)\n",
     "show_structure_3d(cif_content)"
    ]
   },
@@ -331,8 +353,8 @@
    "source": [
     "# Load and plot PAE matrix\n",
     "# Note: Adjust the file path to your actual PAE JSON file\n",
-    "pae_json_file_path = inference_submit_job_response[\"confidences_output_path\"]\n",
-    "pae, token_chain_ids = extract_pae_from_json(pae_json_file_path)\n",
+    "confidence_json_file_path = inference_output_path_info[\"confidences\"]\n",
+    "pae, token_chain_ids = extract_pae_from_json(confidence_json_file_path)\n",
     "plot_pae_matrix(pae,token_chain_ids)"
    ]
   },
@@ -344,7 +366,7 @@
    "source": [
     "# Display summary of confidence metrices\n",
     "# Note: Adjust the file path to your actual PAE Summary JSON file\n",
-    "summary_confidences_file = inference_submit_job_response[\"summary_confidences_output_path\"]\n",
+    "summary_confidences_file = inference_output_path_info[\"summary_confidences\"]\n",
     "summary_data = extract_summary_confidences_obj(summary_confidences_file)\n",
     "\n",
     "# Get chain ID list\n",

examples/science/af3-slurm/examples/simple_ipynb_launcher/adm/slurm_client.py.j2

@@ -27,50 +27,55 @@ def is_valid_path(input_path: str, expected_prefix: str) -> bool:
     return os.path.isabs(input_path) and input_path.startswith(expected_prefix)
 
 
-def sanitize_name(name: str) -> str:
+def sanitise_name(name: str) -> str:
     """Returns sanitised version of the name that can be used as a filename."""
     lower_spaceless_name = name.lower().replace(' ', '_')
     allowed_chars = set(string.ascii_lowercase + string.digits + '_-.')
     return ''.join(l for l in lower_spaceless_name if l in allowed_chars)
 
 
-def retrieve_sanitised_input_file_name(input_file: str) -> Optional[str]:
+def retrieve_sanitised_input_file_name(input_file: str) -> Optional[List[str]]:
     """Retrieves a sanitised version of the input file name."""
     if not os.path.exists(input_file):
         raise FileNotFoundError(f"Input file '{input_file}' does not exist.")
     try:
-        with open(input_file, 'r') as f:
+        with open(input_file, 'r', encoding="utf-8") as f:
             parser = ijson.items(f, '')
             try:
                 root = next(parser)
             except StopIteration:
                 raise ValueError("Empty JSON file.")
 
-            # Case 1: root is a list
+            if not isinstance(root, list) and not isinstance(root, dict):
+                raise ValueError(
+                    "[ERROR] Top-level JSON must be a list or a dictionary.")
+
+            items: Optional[List[Dict]] = None
             if isinstance(root, list):
-                if len(root) != 1:
+                if len(root) <= 0:
                     raise ValueError(
-                        "[ERROR] The AlphaFold Server format is not supported: the list must contain exactly one dictionary item.")
-                item = root[0]
-                if not isinstance(item, dict):
-                    raise ValueError(
-                        "[ERROR] Item in Alphafold Server item must be a dictionary.")
+                        "[ERROR] The AlphaFold Server format is not supported: the list must contain at least one dictionary item.")
+                items = root
 
             # Case 2: root is a dict
             elif isinstance(root, dict):
-                item = root
-            else:
-                raise ValueError(
-                    "[ERROR] Top-level JSON must be a list or a dictionary.")
+                items = [root]
 
             # Validate and extract 'name'
-            name = item.get('name')
-            if not isinstance(name, str):
-                raise ValueError(
-                    "[ERROR] 'name' key must exist and be a string.")
-
-            return sanitize_name(name)
-
+            sanitised_name_list = []
+            for item in items:
+                if not isinstance(item, dict):
+                    raise ValueError(
+                        "Item in Alphafold Format must be a dictionary.")
+                if 'name' not in item:
+                    raise ValueError(
+                        "'name' key must exist in the dictionary.")
+                name = item.get('name')
+                if not isinstance(name, str) or not name:
+                    raise ValueError(
+                        "'name' key must exist, contains value and be a string.")
+                sanitised_name_list.append(sanitise_name(name))
+            return sanitised_name_list
     except FileNotFoundError:
         print(f"[ERROR] Input file '{input_file}' not found.")
         raise FileNotFoundError(
@@ -226,21 +231,6 @@ class AF3SlurmClient:
             output_dir = os.path.join(
                 f"{job_type}_result", file_name, f"{timestamp}")
 
-        inference_model_cif_output_path: Optional[str] = None
-        inference_summary_confidences_output_path: Optional[str] = None
-        inference_confidences_output_path: Optional[str] = None
-
-        sanitised_input_name: Optional[str] = retrieve_sanitised_input_file_name(
-            input_file=input_file
-        )
-
-        inference_model_cif_output_path = os.path.join(
-            output_dir, f"{sanitised_input_name}", f"{sanitised_input_name}_model.cif")
-        inference_summary_confidences_output_path = os.path.join(
-            output_dir, f"{sanitised_input_name}", f"{sanitised_input_name}_summary_confidences.json")
-        inference_confidences_output_path = os.path.join(
-            output_dir, f"{sanitised_input_name}", f"{sanitised_input_name}_confidences.json")
-
         job_config = {
             "account": self.af3_config["service_user"],
             "tasks": 1,
@@ -264,16 +254,31 @@ class AF3SlurmClient:
         response = self.submit_base_job(
             job_config, job_type, script_options, input_file, output_dir)
 
-        response["cif_output_path"] = inference_model_cif_output_path
-        response["summary_confidences_output_path"] = inference_summary_confidences_output_path
-        response["confidences_output_path"] = inference_confidences_output_path
-        print(
-            f"[INFO] Inference model CIF output path : {inference_model_cif_output_path}")
-        print(
-            f"[INFO] Inference confidences output path : {inference_confidences_output_path}")
-        print(
-            f"[INFO] Inference summary confidences output path : {inference_summary_confidences_output_path}")
-
+        sanitised_input_name_list: Optional[List[str]] = retrieve_sanitised_input_file_name(
+            input_file=input_file
+        )
+        output_path_list: List[Dict] = []
+        for sanitised_input_name in sanitised_input_name_list:
+            inference_model_cif_output_path = os.path.join(
+                output_dir, f"{sanitised_input_name}", f"{sanitised_input_name}_model.cif")
+            inference_summary_confidences_output_path = os.path.join(
+                output_dir, f"{sanitised_input_name}", f"{sanitised_input_name}_summary_confidences.json")
+            inference_confidences_output_path = os.path.join(
+                output_dir, f"{sanitised_input_name}", f"{sanitised_input_name}_confidences.json")
+            output_path_list.append({
+                "model_cif": inference_model_cif_output_path,
+                "summary_confidences": inference_summary_confidences_output_path,
+                "confidences": inference_confidences_output_path
+            })
+
+            print(f"[INFO] Output path for input file: {sanitised_input_name}")
+            print(
+                f"         ├── Inference model CIF output path : {inference_model_cif_output_path}")
+            print(
+                f"         ├── Inference confidences output path : {inference_confidences_output_path}")
+            print(
+                f"         └── Inference summary confidences output path : {inference_summary_confidences_output_path}")
+        response["output_path_list"] = output_path_list
         return response
 
     def submit_data_pipeline_job(self,  input_file: str, partition_name: str, output_dir: Optional[str] = None) -> Optional[Dict]:
@@ -291,10 +296,6 @@ class AF3SlurmClient:
                 f"{job_type}_result",  file_name, f"{timestamp}"
             )
 
-        sanitised_input_name = retrieve_sanitised_input_file_name(input_file)
-        datapipeline_output_file_path = os.path.join(
-            output_dir, f"{sanitised_input_name}", f"{sanitised_input_name}_data.json")
-
         job_config = {
             "account": self.af3_config["service_user"],
             "tasks": 1,
@@ -316,7 +317,20 @@ class AF3SlurmClient:
         response = self.submit_base_job(
             job_config, job_type, script_options, input_file, output_dir
         )
-        response["datapipeline_output_path"] = datapipeline_output_file_path
+        sanitised_input_name_list = retrieve_sanitised_input_file_name(
+            input_file)
+        output_path_list: List[Dict] = []
+        for sanitised_input_name in sanitised_input_name_list:
+            datapipeline_output_file_path = os.path.join(
+                output_dir, f"{sanitised_input_name}", f"{sanitised_input_name}_data.json")
+            output_path_list.append({
+                "output_path": datapipeline_output_file_path
+            })
+            print(f"[INFO] Output path for input file: {sanitised_input_name}")
+            print(
+                f"         └── Datapipeline output path : {datapipeline_output_file_path}")
+
+        response["output_path_list"] = output_path_list
         print(
             f"[INFO] Datapipeline output path : {datapipeline_output_file_path}")
         return response