refactor

app.py

@@ -1,3 +1,4 @@
+import uuid
 import gradio as gr
 from pathlib import Path
 from Bio.PDB import MMCIFParser, PDBIO
@@ -35,38 +36,92 @@ def convert_cif_to_pdb(cif_path, pdb_path):
     io.set_structure(structure)
     io.save(pdb_path)
     
-    return pdb_path
 
 
-def predict(sequence: str):
+def predict(sequence: str) -> str:
+    """Predict protein structure from amino acid sequence using Boltz model.
+    
+    Args:
+        sequence (str): Amino acid sequence to predict structure for
+        
+    Returns:
+        str: HTML iframe containing 3D molecular visualization
+    """
+    # Create FASTA file with sequence
     seq_file = Path("sequence.fasta")
-    with open(seq_file, "w") as f:
-        f.write(f">A|protein\n{sequence}")
-    logger.info(f"Predicting {seq_file.stem} with project code {FOLDING_PROJECT_CODE}")
-    output_dir = Path("boltz_results")
+    _write_fasta_file(seq_file, sequence)
+    
+    # Set up unique output directory
+    seq_id = str(uuid.uuid4())
+    output_dir = Path(f"sequence_{seq_id}")
     output_dir.mkdir(parents=True, exist_ok=True)
-    boltz_predict(source=seq_file, project_code=FOLDING_PROJECT_CODE, output=output_dir, unzip=True)
+    
+    # Run Boltz prediction
+    logger.info(f"Predicting {seq_file.stem} with project code {FOLDING_PROJECT_CODE}")
+    boltz_predict(
+        source=seq_file,
+        project_code=FOLDING_PROJECT_CODE, 
+        output=output_dir,
+        unzip=True
+    )
     logger.info("Prediction done. Output directory: %s", output_dir)
-    outs = list(output_dir.rglob("*_model_0.cif"))
-    logger.info("Output files: %s", outs)
-    input_pdb = convert_cif_to_pdb(str(outs[0]), "input.pdb")
-    mol = molecule(
-            str(input_pdb),
-            str(input_pdb),
-            lenSeqs=1,
-            num_res=len(sequence),
-            selectedResidues=list(range(1, len(sequence) + 1)),
-            allSeqs=[sequence],
-            sequences=[{"Score": 0, "RMSD": 0, "Recovery": 0, "Mean pLDDT": 0, "seq": sequence}],
-            random_dir=output_dir,
-        )
-    mol = f"""<iframe style="width: 800px; height: 1300px" name="result" allow="midi; geolocation; microphone; camera; 
-    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
-    allow-scripts allow-same-origin allow-popups 
-    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
-    allowpaymentrequest="" frameborder="0" srcdoc='{mol}'></iframe>"""
-
-    return mol
+
+    # Convert output CIF to PDB
+    pred_cif = list(output_dir.rglob("*_model_0.cif"))[0]
+    logger.info("Output file: %s", pred_cif)
+    
+    converted_pdb_path = output_dir / "pred.pdb"
+    convert_cif_to_pdb(str(pred_cif), str(converted_pdb_path))
+    logger.info("Converted PDB file: %s", converted_pdb_path)
+
+    # Generate molecular visualization
+    mol = _create_molecule_visualization(
+        converted_pdb_path,
+        sequence,
+        output_dir
+    )
+    
+    return _wrap_in_iframe(mol)
+
+
+def _write_fasta_file(filepath: Path, sequence: str) -> None:
+    """Write sequence to FASTA file."""
+    with open(filepath, "w") as f:
+        f.write(f">A|protein\n{sequence}")
+
+
+def _create_molecule_visualization(pdb_path: Path, sequence: str, output_dir: Path) -> str:
+    """Create molecular visualization using molecule module."""
+    return molecule(
+        str(pdb_path),
+        str(pdb_path),
+        lenSeqs=1,
+        num_res=len(sequence),
+        selectedResidues=list(range(1, len(sequence) + 1)),
+        allSeqs=[sequence],
+        sequences=[{
+            "Score": 0,
+            "RMSD": 0, 
+            "Recovery": 0,
+            "Mean pLDDT": 0,
+            "seq": sequence
+        }],
+        random_dir=output_dir
+    )
+
+
+def _wrap_in_iframe(content: str) -> str:
+    """Wrap content in an HTML iframe with appropriate styling and permissions."""
+    return f"""<iframe 
+        style="width: 800px; height: 1300px" 
+        name="result" 
+        allow="midi; geolocation; microphone; camera; display-capture; encrypted-media;"
+        sandbox="allow-modals allow-forms allow-scripts allow-same-origin allow-popups allow-top-navigation-by-user-activation allow-downloads"
+        allowfullscreen=""
+        allowpaymentrequest=""
+        frameborder="0"
+        srcdoc='{content}'
+    ></iframe>"""
 
 demo = gr.Blocks()
 
@@ -74,7 +129,7 @@ with demo:
     gr.Markdown("# Input")
     with gr.Row():
         with gr.Column():
-            sequence = gr.Textbox(label="Sequence", value="MALWMRLLPLLALLALWGPDPAAA")
+            sequence = gr.Textbox(label="Sequence", value="")
     gr.Markdown("# Output")
     with gr.Row():
         predict_btn = gr.Button("Predict")