fix api calls

app.py

@@ -4,27 +4,29 @@ import io
 import numpy as np
 import os
 import traceback
-# import spaces
 from esm.sdk import client
-from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig
+from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig, ESMProteinError
 from esm.utils.structure.protein_chain import ProteinChain
 from Bio.Data import PDBData
 import biotite.structure as bs
 from biotite.structure.io import pdb
 from esm.utils import residue_constants as RC
-
+import requests
 from dotenv import load_dotenv
+import torch
 
 load_dotenv()
-# Initialize the model
-token = os.environ.get("ESM_API_TOKEN")
-if not token:
+
+API_URL = "https://forge.evolutionaryscale.ai/api/v1"
+MODEL = "esm3-open-2024-03"
+API_TOKEN = os.environ.get("ESM_API_TOKEN")
+if not API_TOKEN:
     raise ValueError("ESM_API_TOKEN environment variable is not set")
 
 model = client(
-    model="esm3-medium-2024-03",
-    url="https://forge.evolutionaryscale.ai",
-    token=token,
+    model=MODEL,
+    url=API_URL,
+    token="2x0lifRJCpo8klurAJtRom"
 )
 
 amino3to1 = {
@@ -34,6 +36,7 @@ amino3to1 = {
     'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y'
 }
 
+
 def read_pdb_io(pdb_file):
     if isinstance(pdb_file, io.StringIO):
         pdb_content = pdb_file.getvalue()
@@ -129,15 +132,19 @@ def add_noise_to_coordinates(protein: ESMProtein, noise_level: float) -> ESMProt
 def run_structure_prediction(protein: ESMProtein) -> ESMProtein:
     structure_prediction_config = GenerationConfig(
         track="structure",
-        num_steps=40,
+        num_steps=10,
         temperature=0.7,
     )
     try:
-        structure_prediction = model.generate(protein, structure_prediction_config)
-        if isinstance(structure_prediction, ESMProtein):
-            return structure_prediction
+        response = model.generate(protein, structure_prediction_config)
+        
+        if isinstance(response, ESMProtein):
+            return response
+        elif isinstance(response, ESMProteinError):
+            print(f"ESMProteinError during structure prediction: {response.error_msg}")
+            return None
         else:
-            raise ValueError(f"Unexpected result type: {type(structure_prediction)}")
+            raise ValueError(f"Unexpected response type: {type(response)}")
     except Exception as e:
         print(f"Error during structure prediction: {str(e)}")
         return None
@@ -149,15 +156,55 @@ def align_after_prediction(protein: ESMProtein, structure_prediction: ESMProtein
     try:
         structure_prediction_chain = structure_prediction.to_protein_chain()
         protein_chain = protein.to_protein_chain()
-        structure_indices = np.arange(0, len(structure_prediction_chain.sequence))
-        aligned_chain = structure_prediction_chain.align(protein_chain, mobile_inds=structure_indices, target_inds=structure_indices)
-        crmsd = structure_prediction_chain.rmsd(protein_chain, mobile_inds=structure_indices, target_inds=structure_indices)    
+        
+        # Ensure both chains have the same length
+        min_length = min(len(structure_prediction_chain.sequence), len(protein_chain.sequence))
+        structure_indices = np.arange(0, min_length)
+        
+        # Perform alignment
+        aligned_chain = structure_prediction_chain.align(
+            protein_chain, 
+            mobile_inds=structure_indices, 
+            target_inds=structure_indices
+        )
+        
+        # Calculate RMSD
+        crmsd = structure_prediction_chain.rmsd(
+            protein_chain, 
+            mobile_inds=structure_indices, 
+            target_inds=structure_indices
+        )
+        
         return ESMProtein.from_protein_chain(aligned_chain), crmsd
     except AttributeError as e:
         print(f"Error during alignment: {str(e)}")
         print(f"Structure prediction type: {type(structure_prediction)}")
         print(f"Structure prediction attributes: {dir(structure_prediction)}")
         return None, float('inf')
+    except Exception as e:
+        print(f"Unexpected error during alignment: {str(e)}")
+        return None, float('inf')
+    
+def visualize_after_pred(protein: ESMProtein, aligned: ESMProtein):
+    if aligned is None:
+        return py3Dmol.view(width=800, height=600)
+    
+    view = py3Dmol.view(width=800, height=600)
+    view.addModel(protein_to_pdb(protein), "pdb")
+    view.setStyle({"cartoon": {"color": "lightgrey"}})
+    view.addModel(protein_to_pdb(aligned), "pdb")
+    view.setStyle({"model": 1}, {"cartoon": {"color": "lightgreen"}})
+    view.zoomTo()
+    return view
+
+def protein_to_pdb(protein: ESMProtein):
+    pdb_str = ""
+    for i, (aa, coords) in enumerate(zip(protein.sequence, protein.coordinates)):
+        for j, atom in enumerate(RC.atom_types):
+            if not torch.isnan(coords[j][0]):
+                x, y, z = coords[j].tolist()
+                pdb_str += f"ATOM  {i*37+j+1:5d}  {atom:3s} {aa:3s} A{i+1:4d}    {x:8.3f}{y:8.3f}{z:8.3f}\n"
+    return pdb_str
 
 def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_frames: int):
     protein = get_protein(pdb_file)
@@ -168,30 +215,18 @@ def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_fr
         
         for i in range(num_runs):
             structure_prediction = run_structure_prediction(noisy_protein)
-            aligned, crmsd = align_after_prediction(protein, structure_prediction)
-            if aligned is not None:
-                runs.append((crmsd, aligned))
+            if structure_prediction is not None:
+                aligned, crmsd = align_after_prediction(protein, structure_prediction)
+                if aligned is not None:
+                    runs.append((crmsd, aligned))
     
     if not runs:
         return None, "No successful predictions"
     
-    best_aligned = sorted(runs)[0]
+    best_aligned = sorted(runs, key=lambda x: x[0])[0]
     view = visualize_after_pred(protein, best_aligned[1])
     return view, f"Best cRMSD: {best_aligned[0]:.4f}"
 
-def visualize_after_pred(protein: ESMProtein, aligned: ESMProtein):
-    if aligned is None:
-        return py3Dmol.view(width=800, height=600)
-    
-    view = py3Dmol.view(width=800, height=600)
-    view.addModel(protein.to_pdb_string(), "pdb")
-    view.setStyle({"cartoon": {"color": "lightgrey"}})
-    view.addModel(aligned.to_pdb_string(), "pdb")
-    view.setStyle({"model": 1}, {"cartoon": {"color": "lightgreen"}})
-    view.zoomTo()
-    return view
-
-# @spaces.GPU()
 def run_prediction(pdb_file, num_runs, noise_level, num_frames):
     try:
         if pdb_file is None:
@@ -199,7 +234,7 @@ def run_prediction(pdb_file, num_runs, noise_level, num_frames):
         
         view, crmsd_text = prediction_visualization(pdb_file, num_runs, noise_level, num_frames)
         if view is None:
-            return "No successful predictions were made. Try adjusting the parameters.", crmsd_text
+            return "No successful predictions were made. Try adjusting the parameters or check the PDB file.", crmsd_text
         
         html = view._make_html()
         return f"""