Files
D-SCRIPT/dscript/tests/test_language_model.py
2022-08-18 12:10:32 -04:00

36 lines
872 B
Python

import os
import shutil
import subprocess as sp
from Bio import SeqIO
from dscript.language_model import (
lm_embed,
embed_from_fasta,
)
class TestLanguageModel:
@classmethod
def setup_class(cls):
cmd = "python setup.py install"
proc = sp.Popen(cmd.split())
proc.wait()
os.makedirs("./tmp-dscript-testing/", exist_ok=True)
@classmethod
def teardown_class(cls):
shutil.rmtree("./tmp-dscript-testing/")
def test_lm_embed(self):
seqs = list(SeqIO.parse("dscript/tests/test.fasta", "fasta"))
for seqrec in seqs:
x = lm_embed(str(seqrec.seq))
assert x.shape[1] == len(seqrec.seq)
def embed_from_fasta(self):
embed_from_fasta(
"dscript/tests/test.fasta",
"tmp-dscript-testing/test_embed.h5",
verbose=True,
)