fix inference_app.py

inference_app.py

@@ -106,7 +106,7 @@ from Bio import PDB
 from Bio.PDB.PDBIO import PDBIO
 
 def extract_coordinates_from_pdb(filename):
-    """
+    r"""
     Extracts atom coordinates from a PDB file and returns them as a list of tuples.
     Each tuple contains (x, y, z) coordinates of an atom.
     """
@@ -288,7 +288,7 @@ class PairedPDB(HeteroData):  # type: ignore
       
 #create_graph takes inputs apo_ligand, apo_residue and paired holo as pdb3(ground truth). 
 def create_graph(pdb1, pdb2, pdb3='/home/sukanya/iitm_bisect_pinder_submission/test_out.pdb', k=5):
-    """
+    r"""
     Create a heterogeneous graph from two PDB files, with the ligand and receptor
     as separate nodes, and their respective features and edges.
     
@@ -336,7 +336,7 @@ def create_graph(pdb1, pdb2, pdb3='/home/sukanya/iitm_bisect_pinder_submission/t
 
 
 def update_pdb_coordinates_from_tensor(input_filename, output_filename, coordinates_tensor):
-    """
+    r"""
     Updates atom coordinates in a PDB file with new transformed coordinates provided in a tensor.
 
     Parameters:
@@ -400,7 +400,7 @@ def update_pdb_coordinates_from_tensor(input_filename, output_filename, coordina
     return output_filename
  
 def merge_pdb_files(file1, file2, output_file):
-    """
+    r"""
     Merges two PDB files by concatenating them without altering their contents.
     
     Parameters:
@@ -421,7 +421,7 @@ def merge_pdb_files(file1, file2, output_file):
     
 class MPNNLayer(MessagePassing):
     def __init__(self, emb_dim=64, edge_dim=4, aggr='add'):
-        """Message Passing Neural Network Layer
+        r"""Message Passing Neural Network Layer
 
         Args:
             emb_dim: (int) - hidden dimension d
@@ -451,7 +451,7 @@ class MPNNLayer(MessagePassing):
           )
 
     def forward(self, h, edge_index, edge_attr):
-        """
+        r"""
         The forward pass updates node features h via one round of message passing.
 
         As our MPNNLayer class inherits from the PyG MessagePassing parent class,
@@ -474,7 +474,7 @@ class MPNNLayer(MessagePassing):
         return out
 
     def message(self, h_i, h_j, edge_attr):
-        """Step (1) Message
+        r"""Step (1) Message
 
         The message() function constructs messages from source nodes j 
         to destination nodes i for each edge (i, j) in edge_index.
@@ -502,7 +502,7 @@ class MPNNLayer(MessagePassing):
         return self.mlp_msg(msg)
     
     def aggregate(self, inputs, index):
-        """Step (2) Aggregate
+        r"""Step (2) Aggregate
 
         The aggregate function aggregates the messages from neighboring nodes,
         according to the chosen aggregation function ('sum' by default).
@@ -517,7 +517,7 @@ class MPNNLayer(MessagePassing):
         return scatter(inputs, index, dim=self.node_dim, reduce=self.aggr)
     
     def update(self, aggr_out, h):
-        """
+        r"""
         Step (3) Update
 
         The update() function computes the final node features by combining the 
@@ -541,7 +541,7 @@ class MPNNLayer(MessagePassing):
         return (f'{self.__class__.__name__}(emb_dim={self.emb_dim}, aggr={self.aggr})')
 class MPNNModel(Module):
     def __init__(self, num_layers=4, emb_dim=64, in_dim=11, edge_dim=4, out_dim=1):
-        """Message Passing Neural Network model for graph property prediction
+        r"""Message Passing Neural Network model for graph property prediction
 
         Args:
             num_layers: (int) - number of message passing layers L
@@ -570,7 +570,7 @@ class MPNNModel(Module):
         self.lin_pred = Linear(emb_dim, out_dim)
         
     def forward(self, data):
-        """
+        r"""
         Args:
             data: (PyG.Data) - batch of PyG graphs
 
@@ -592,7 +592,7 @@ class MPNNModel(Module):
 
 class EquivariantMPNNLayer(MessagePassing):
     def __init__(self, emb_dim=64,  aggr='add'):
-        """Message Passing Neural Network Layer
+        r"""Message Passing Neural Network Layer
 
         This layer is equivariant to 3D rotations and translations.
 
@@ -630,7 +630,7 @@ class EquivariantMPNNLayer(MessagePassing):
         # ===========================================
 
     def forward(self, h, pos, edge_index):
-        """
+        r"""
         The forward pass updates node features h via one round of message passing.
 
         Args:
@@ -667,7 +667,7 @@ class EquivariantMPNNLayer(MessagePassing):
     #   ...  
     #
     def aggregate(self, inputs, index):
-        """The aggregate function aggregates the messages from neighboring nodes,
+        r"""The aggregate function aggregates the messages from neighboring nodes,
         according to the chosen aggregation function ('sum' by default).
 
         Args:
@@ -701,7 +701,7 @@ class EquivariantMPNNLayer(MessagePassing):
     
 class FinalMPNNModel(MPNNModel):
     def __init__(self, num_layers=4, emb_dim=64, in_dim=3,  num_heads = 2):
-        """Message Passing Neural Network model for graph property prediction
+        r"""Message Passing Neural Network model for graph property prediction
 
         This model uses both node features and coordinates as inputs, and
         is invariant to 3D rotations and translations (the constituent MPNN layers
@@ -734,7 +734,7 @@ class FinalMPNNModel(MPNNModel):
         # self.pool = global_mean_pool
 
     def naive_single(self, receptor, ligand , receptor_edge_index , ligand_edge_index):
-        """
+        r"""
         Processes a single receptor-ligand pair.
 
         Args:
@@ -775,7 +775,7 @@ class FinalMPNNModel(MPNNModel):
    
 
     def forward(self, data):
-        """
+        r"""
         The main forward pass of the model.
 
         Args: