Update app.py

app.py

@@ -176,40 +176,112 @@ def parse_json_input(json_data: List[Dict]) -> Dict:
                 })
     return components
 
+
 def create_protenix_json(input_data: Dict) -> List[Dict]:
-    """Convert UI inputs to Protenix JSON format"""
     sequences = []
     
-    for pc in input_data["protein_chains"]:
-        sequences.append({
-            "proteinChain": {
-                "sequence": pc["sequence"],
-                "count": pc["count"]
-            }
-        })
+    # Process protein chains
+    for pc in input_data.get("protein_chains", []):
+        # Check that the row has both columns and the sequence is nonempty.
+        if len(pc) >= 2 and pc[0].strip():
+            sequences.append({
+                "proteinChain": {
+                    "sequence": pc[0].strip(),
+                    "count": int(pc[1]) if pc[1] else 1
+                }
+            })
     
-    for dna in input_data["dna_sequences"]:
-        sequences.append({
-            "dnaSequence": {
-                "sequence": dna["sequence"],
-                "count": dna["count"]
-            }
-        })
+    # Process DNA sequences
+    for dna in input_data.get("dna_sequences", []):
+        if len(dna) >= 2 and dna[0].strip():
+            sequences.append({
+                "dnaSequence": {
+                    "sequence": dna[0].strip(),
+                    "count": int(dna[1]) if dna[1] else 1
+                }
+            })
+
+    # Process RNA sequences
+    for rna in input_data.get("rna_sequences", []):
+        if len(rna) >= 2 and rna[0].strip():
+            sequences.append({
+                "rnaSequence": {
+                    "sequence": rna[0].strip(),
+                    "count": int(rna[1]) if rna[1] else 1
+                }
+            })
     
-    for lig in input_data["ligands"]:
-        sequences.append({
-            "ligand": {
-                "ligand": lig["type"],
-                "count": lig["count"]
-            }
-        })
+    # Process ligands
+    for lig in input_data.get("ligands", []):
+        if len(lig) >= 2 and lig[0].strip():
+            sequences.append({
+                "ligand": {
+                    "ligand": lig[0].strip(),
+                    "count": int(lig[1]) if lig[1] else 1
+                }
+            })
     
     return [{
         "sequences": sequences,
-        "name": input_data["complex_name"]
+        "name": input_data.get("complex_name")+f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:3]}"
     }]
 
 
+def update_json(complex_name, protein_chains, dna_sequences, rna_sequences, ligands):
+    sequences_list = []
+    
+    # Process protein chains (DataFrame with headers: ["Sequence", "Count"])
+    if protein_chains:
+        for row in protein_chains:
+            # Check if the row is valid and non-empty
+            if row and len(row) >= 2 and row[0]:
+                sequences_list.append({
+                    "proteinChain": {
+                        "sequence": row[0],
+                        "count": row[1]
+                    }
+                })
+                
+    # Process DNA sequences
+    if dna_sequences:
+        for row in dna_sequences:
+            if row and len(row) >= 2 and row[0]:
+                sequences_list.append({
+                    "dnaSequence": {
+                        "sequence": row[0],
+                        "count": row[1]
+                    }
+                })
+                
+    # Process RNA sequences
+    if rna_sequences:
+        for row in rna_sequences:
+            if row and len(row) >= 2 and row[0]:
+                sequences_list.append({
+                    "rnaSequence": {
+                        "sequence": row[0],
+                        "count": row[1]
+                    }
+                })
+    
+    # Process ligands (DataFrame with headers: ["Ligand Type", "Count"])
+    if ligands:
+        for row in ligands:
+            if row and len(row) >= 2 and row[0]:
+                sequences_list.append({
+                    "ligand": {
+                        "ligand": row[0],
+                        "count": row[1]
+                    }
+                })
+    
+    return {
+        "sequences": sequences_list,
+        "name": complex_name
+    }
+
+
+
 #@torch.inference_mode()
 @spaces.GPU(duration=120)  # Specify a duration to avoid timeout
 def predict_structure(input_collector: dict):
@@ -225,7 +297,7 @@ def predict_structure(input_collector: dict):
         print(input_collector)
 
         # Handle JSON input
-        if input_collector["json"]:
+        if "json" in input_collector:
             # Handle different input types
             if isinstance(input_collector["json"], str):  # Example JSON case (file path)
                 input_data = json.load(open(input_collector["json"]))
@@ -406,31 +478,44 @@ with gr.Blocks(title="FoldMark", css=custom_css) as demo:
                 headers=["Sequence", "Count"],
                 datatype=["str", "number"],
                 row_count=1,
-                col_count=(2, "fixed")
+                col_count=(2, "fixed"),
+                type="array"
             )
             
         # Repeat for other groups
-        with gr.Accordion(label="DNA Sequences", open=True):
+        with gr.Accordion(label="DNA Sequences (A T G C)", open=True):
             dna_sequences = gr.Dataframe(
                 headers=["Sequence", "Count"],
                 datatype=["str", "number"],
-                row_count=1
+                row_count=1,
+                type="array"
+            )
+
+        with gr.Accordion(label="RNA Sequences (A U G C)", open=True):
+            rna_sequences = gr.Dataframe(
+                headers=["Sequence", "Count"],
+                datatype=["str", "number"],
+                row_count=1,
+                type="array"
             )
             
         with gr.Accordion(label="Ligands", open=True):
             ligands = gr.Dataframe(
                 headers=["Ligand Type", "Count"],
                 datatype=["str", "number"],
-                row_count=1
+                row_count=1,
+                type="array"
             )
             
         manual_output = gr.JSON(label="Generated JSON")
         
-        complex_name.change(
-            fn=lambda x: {"complex_name": x},
-            inputs=complex_name,
-            outputs=manual_output
-        )
+        # Attach a change event to all widgets so that any change updates the JSON output.
+        for widget in [complex_name, protein_chains, dna_sequences, rna_sequences, ligands]:
+            widget.change(
+                fn=update_json,
+                inputs=[complex_name, protein_chains, dna_sequences, rna_sequences, ligands],
+                outputs=manual_output
+            )
 
         # Shared prediction components
         with gr.Row():
@@ -450,8 +535,8 @@ with gr.Blocks(title="FoldMark", css=custom_css) as demo:
 
         # Map inputs to a dictionary
         submit_btn.click(
-            fn=lambda c, p, d, l, w: {"data": {"complex_name": c, "protein_chains": p, "dna_sequences": d, "ligands": l}, "watermark": w},
-            inputs=[complex_name, protein_chains, dna_sequences, ligands, add_watermark1],
+            fn=lambda c, p, d, r, l, w: {"data": {"complex_name": c, "protein_chains": p, "dna_sequences": d, "rna_sequences": r, "ligands": l}, "watermark": w},
+            inputs=[complex_name, protein_chains, dna_sequences, rna_sequences, ligands, add_watermark1],
             outputs=input_collector
         ).then(
             fn=predict_structure,