File size: 559 Bytes
f8402f9
87c0dbc
f8402f9
87c0dbc
f8402f9
 
 
 
 
87c0dbc
f8402f9
 
87c0dbc
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from Bio.PDB.Structure import Structure
from transformers import T5EncoderModel, T5Tokenizer

from protention.attention import get_protT5, get_structure


def test_get_structure():
    pdb_id = "1AKE"
    structure = get_structure(pdb_id)

    assert structure is not None
    assert isinstance(structure, Structure)


def test_get_protT5():
    result = get_protT5()

    assert result is not None
    assert isinstance(result, tuple)

    tokenizer, model = result

    assert isinstance(tokenizer, T5Tokenizer)
    assert isinstance(model, T5EncoderModel)