File size: 660 Bytes
eb9ae1f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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

from transformers import (GPT2LMHeadModel, GPT2TokenizerFast, T5EncoderModel,
                          T5Tokenizer)

from hexviz.models import get_protT5, get_zymctrl


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_zymctrl():
    result = get_zymctrl()

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

    tokenizer, model = result

    assert isinstance(tokenizer, GPT2TokenizerFast)
    assert isinstance(model, GPT2LMHeadModel)