add jsons

app.py

@@ -1,4 +1,5 @@
 import functools
+import json
 import os
 import tempfile
 import time
@@ -9,6 +10,7 @@ import pandas as pd
 from loguru import logger
 
 from stoic.model import Stoic
+from stoic.predict_stoichiometry import _build_af3_input_json
 
 MAX_CHAINS = 26
 
@@ -73,6 +75,8 @@ def predict(sequences_text: str, top_n: int, return_weights: bool):
     stoich_md = table + "\n".join(legend_lines)
     stoich_csv_path = _save_csv(pd.DataFrame(stoich_csv_rows), "stoichiometry_results.csv")
 
+    af3_json_paths = _save_af3_jsons(results)
+
     plot_updates = [gr.update(value=None, visible=False)] * MAX_CHAINS
     weights_csv_update = gr.update(value=None, visible=False)
 
@@ -90,6 +94,7 @@ def predict(sequences_text: str, top_n: int, return_weights: bool):
         stoich_md,
         f"{elapsed:.2f}s",
         gr.update(value=stoich_csv_path, visible=True),
+        gr.update(value=af3_json_paths, visible=True),
         weights_csv_update,
         *plot_updates,
     )
@@ -101,6 +106,18 @@ def _save_csv(df: pd.DataFrame, filename: str) -> str:
     return path
 
 
+def _save_af3_jsons(results: list[dict]) -> list[str]:
+    """Generate AF3-style JSON files for each stoichiometry candidate."""
+    paths = []
+    for rank, candidate in enumerate(results, 1):
+        af3_json = _build_af3_input_json(f"stoic_rank{rank}", [candidate])
+        path = os.path.join(tempfile.gettempdir(), f"stoic_rank{rank}_af3.json")
+        with open(path, "w") as f:
+            json.dump(af3_json, f, indent=2)
+        paths.append(path)
+    return paths
+
+
 def build_chain_dfs(residue_predictions, chain_labels):
     pred_residues = residue_predictions["pred_residues"]
     attention_mask = residue_predictions["attention_mask"]
@@ -158,6 +175,11 @@ with gr.Blocks(title="Stoic - Protein Stoichiometry Prediction") as app:
             label="Download Stoichiometry Results (CSV)",
             visible=False,
         )
+        af3_json_download = gr.File(
+            label="Download AF3 Input JSON(s)",
+            file_count="multiple",
+            visible=False,
+        )
         weights_csv_download = gr.File(
             label="Download Residue Weights (CSV)",
             visible=False,
@@ -186,6 +208,7 @@ with gr.Blocks(title="Stoic - Protein Stoichiometry Prediction") as app:
             results_output,
             run_time,
             stoich_csv_download,
+            af3_json_download,
             weights_csv_download,
             *chain_plots,
         ],