Spaces:
Sleeping
Sleeping
File size: 1,083 Bytes
466a8f2 f8402f9 87c0dbc f8402f9 ebbe380 466a8f2 f8402f9 87c0dbc f8402f9 87c0dbc b7ab123 87c0dbc 466a8f2 ebbe380 466a8f2 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import torch
from Bio.PDB.Structure import Structure
from transformers import T5EncoderModel, T5Tokenizer
from protention.attention import (ModelType, get_attention, get_protT5,
get_sequences, 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_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()
assert result is not None
assert isinstance(result, tuple)
tokenizer, model = result
assert isinstance(tokenizer, T5Tokenizer)
assert isinstance(model, T5EncoderModel)
def test_get_attention_tape():
result = get_attention("1AKE", model=ModelType.tape_bert)
assert result is not None
assert result.shape == torch.Size([12,12,456,456])
|