from transformers import AutoModel, AlbertTokenizer, AutoConfig
import re

def convert_sequence_to_input(sequence: str, model_name = "Rostlab/prot_albert"):
    seq = " ".join([aa for aa in sequence])
    seq = re.sub(r"[UZOB]", "X", seq)
    tokenizer = AlbertTokenizer.from_pretrained(model_name, trust_remote_code=True, do_lower_case=False)
    params = dict(return_tensors="pt", padding="max_length",
                  max_length=128,
                  truncation=True,)
    x = tokenizer(seq, **params)
    return x

def convert_pfam_idx_to_class(pfam_idx: int) -> str:
    """
    Convert the prediction of the model to the corresponding class.
    :param pfam_idx: index of the pfam class
    :return: the Pfam family
    """
    conversion = {"0": "Methyltransf_25", "1": "LRR_1", "2": "Acetyltransf_7", "3": "His_kinase",
     "4": "Bac_transf", "5": "Lum_binding", "6": "DNA_binding_1", "7": "Chromate_transp",
     "8": "Lipase_GDSL_2", "9": "DnaJ_CXXCXGXG"}
    return conversion.get(str(pfam_idx), "Unknown")


model_name = "sayby/prot_albert_pfam"
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)



sequence = "ILDVGTGTGKLESLAEFKRDFIGLDVTKEMMALNRNKGKLLLASATQMPIKDGTFDAIVSSFVLRNLPSTKGYFSEGFRTLKEGG"
x = convert_sequence_to_input(sequence)
output = model(x)
pfam_idx = output["logits"].argmax(dim=-1).item()
pfam = convert_pfam_idx_to_class(pfam_idx)
print(f"The Pfam family is: {pfam}")