SaNano

from transformers import EsmForMaskedLM, EsmTokenizer
import torch

def add_structure_masking(sequence):
    return "#".join(sequence) + "#"

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = EsmForMaskedLM.from_pretrained("novonordisk-red/SaNano").to(device)
tokenizer = EsmTokenizer.from_pretrained("novonordisk-red/SaNano")

sequences = [
    "EVQLVESGGGLVQAGGSLRLSCAASGFTFPTYAMAWFRQAPGKGREFV",
    "QVQLQESGGGLVQAGGSLRLSCAASGRTFSSYAMGWFRQAPGKEREFVAAI",
    "EVQLVESGGGLVQPGGSLRLSCAASGFTFSSYAMSWVRQAPGKG",
    "DVQLVESGGGLVQAGGSLRLSCAASGRTFSTYAMGWFRQAPGKGREFVAGISWS"
]

# To convert to structure-aware format with masking. This is necessary if structural tokens are not computed
structure_aware_sequences = [add_structure_masking(seq) for seq in sequences]
# ["E#V#Q#L#V#E#S#G#G#G#L#V#Q#A#G#G#S#L#R#L#S#C#A#A#S#G#F#T#F#P#T#Y#A#M#A#W#F#R#Q#A#P#G#...", ...]

# Move input to device and ensure hidden states are returned
inputs = tokenizer(structure_aware_sequences, return_tensors="pt", padding=True)
inputs = {key: value.to(device) for key, value in inputs.items()}
inputs['output_hidden_states'] = True

outputs = model(**inputs)