Add get_sequence

protention/attention.py

@@ -2,7 +2,7 @@ from io import StringIO
 from urllib import request
 
 import torch
-from Bio.PDB import PDBParser, Structure
+from Bio.PDB import PDBParser, Polypeptide, Structure
 from transformers import T5EncoderModel, T5Tokenizer
 
 
@@ -17,6 +17,20 @@ def get_structure(pdb_code: str) -> Structure:
     structure = parser.get_structure(pdb_code, file)
     return structure
 
+def get_sequences(structure: Structure) -> list[str]:
+    """
+    Get list of sequences with residues on a single letter format
+
+    Residues not in the standard 20 amino acids are replaced with X
+    """
+    sequences = []
+    for seq in structure.get_chains():
+        residues = [residue.get_resname() for residue in seq.get_residues()]
+        # TODO ask if using protein_letters_3to1_extended makes sense
+        residues_single_letter = map(lambda x: Polypeptide.protein_letters_3to1.get(x, "X"), residues)
+
+        sequences.append(list(residues_single_letter))
+    return sequences
 
 def get_protT5() -> tuple[T5Tokenizer, T5EncoderModel]:
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -41,11 +55,14 @@ def get_attention(
     """
     # fetch structure
     structure = get_structure(pdb_code)
+    # Get list of sequences
+    sequences = get_sequences(structure)
 
     # get model
     tokenizer, model = get_protT5()
 
     # call model
+    ## Get sequence
 
     # get attention
 

tests/test_attention.py

@@ -1,7 +1,7 @@
 from Bio.PDB.Structure import Structure
 from transformers import T5EncoderModel, T5Tokenizer
 
-from protention.attention import get_protT5, get_structure
+from protention.attention import get_protT5, get_sequences, get_structure
 
 
 def test_get_structure():
@@ -11,6 +11,17 @@ def test_get_structure():
     assert structure is not None
     assert isinstance(structure, Structure)
 
+def test_get_sequences():
+    pdb_id = "1AKE"
+    structure = get_structure(pdb_id)
+    
+    sequences = get_sequences(structure)
+
+    assert sequences is not None
+    assert len(sequences) == 2
+
+    A, B = sequences
+    assert A[:3] == ["M", "R", "I"]
 
 def test_get_protT5():
     result = get_protT5()