from transformers import AutoTokenizer, AutoModel
import logging
import torch

# Suppress warnings about newly initialized 'esm.pooler.dense.bias', 'esm.pooler.dense.weight' layers - these are not used to extract embeddings
logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load the tokenizer and model
model_name = "ChatterjeeLab/FusOn-pLM" 
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
model.to(device)
model.eval()

# Example fusion oncoprotein sequence: MLLT10:PICALM, associated with Acute Myeloid Leukemia (LAML)  
# Amino acids 1-80 are derived from the head gene, MLLT10
# Amino acids 81-119 are derived from the tail gene, PICALM
sequence = "MVSSDRPVSLEDEVSHSMKEMIGGCCVCSDERGWAENPLVYCDGHGCSVAVHQACYGIVQVPTGPWFCRKCESQERAARVPPQMGSVPVMTQPTLIYSQPVMRPPNPFGPVSGAQIQFM"

# Tokenize the input sequence
inputs = tokenizer(sequence, return_tensors="pt", padding=True, truncation=True,max_length=2000)
inputs = {k: v.to(device) for k, v in inputs.items()}

# Get the embeddings
with torch.no_grad():
    outputs = model(**inputs)
    # The embeddings are in the last_hidden_state tensor
    embeddings = outputs.last_hidden_state
    # remove extra dimension
    embeddings = embeddings.squeeze(0)
    # remove BOS and EOS tokens
    embeddings = embeddings[1:-1, :]

# Convert embeddings to numpy array (if needed)
embeddings = embeddings.cpu().numpy()

print("Per-residue embeddings shape:", embeddings.shape)