predict_seq.py

import tempfile
from pathlib import Path
from warnings import filterwarnings
from tmbed.tmbed.run_tmbed import predict, load_models

def predict_tmbed(sequence: str, models, device):
    with tempfile.TemporaryDirectory() as td:
        input_path = Path(td) / "input.fasta"
        output_path = Path(td) / "output.fasta"
        input_path.write_text(f"> Protein Chain\n{sequence}")
        predictions = predict(fasta_file=input_path,
                              output_file=output_path,
                              out_format='3',
                              models=models,
                              device=device)
        pTMBED_prop_5c = list(predictions.values())[0][1]
        tm_prop = pTMBED_prop_5c[:, :2].sum(axis=1)
        sig_prop = pTMBED_prop_5c[:, 2]
        sp_prop = pTMBED_prop_5c[:, 3:].sum(axis=1)
        return tm_prop, sig_prop, sp_prop