Views
No views yet
1
2from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM
3import torch
4
5model_id = "gustoudu81/BerTeleo"
6
7tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
8model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
9
10device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11model = model.to(device).eval()
12
13inputs = tokenizer("ACGTACGTACGT", return_tensors="pt")
14inputs = {k: v.to(device) for k, v in inputs.items()}
15
16with torch.no_grad():
17 hidden_states = model(**inputs)[0]
18
19
20# embedding with mean pooling
21embedding_mean = torch.mean(hidden_states[0], dim=0)
22print(embedding_mean.shape) # expect to be 768
23
24# embedding with max pooling
25embedding_max = torch.max(hidden_states[0], dim=0)[0]
26print(embedding_max.shape) # expect to be 768