Views
No views yet
1from transformers import AutoTokenizer, AutoModel
2import torch
3
4# --- Config ---
5model_name = "InstaDeepAI/IDP-ESM2-150M"
6
7# --- Load model and tokenizer ---
8tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
9model = AutoModel.from_pretrained(model_name)
10model.eval()
11
12# (optional) use GPU if available
13device = "cuda" if torch.cuda.is_available() else "cpu"
14model.to(device)
15
16# --- Input sequences ---
17sequences = [
18 "MDDNHYPHHHHNHHNHHSTSGGCGESQFTTKLSVNTFARTHPMIQNDLIDLDLISGSAFTMKSKSQQ",
19 "PADRDLSSPFGSTVPGVGPNAAAASNAAAAAAAAATAGSNKHQTPPTTFR",
20]
21
22# --- Tokenize ---
23inputs = tokenizer(
24 sequences,
25 return_tensors="pt",
26 padding=True,
27 truncation=True,
28)
29inputs = {k: v.to(device) for k, v in inputs.items()}
30
31# --- Forward pass ---
32with torch.no_grad():
33 outputs = model(**inputs)
34 embeddings = outputs.last_hidden_state # shape: (batch, seq_len, hidden_dim)
35