Views
No views yet
1from transformers import AutoModelForCausalLM
2from tokenizers import Tokenizer
3import torch
4import torch.nn.functional as F
5
6# load model and tokenizer
7model = AutoModelForCausalLM.from_pretrained("hugohrban/progen2-base", trust_remote_code=True)
8tokenizer = Tokenizer.from_pretrained("hugohrban/progen2-base")
9tokenizer.no_padding()
10
11# prepare input
12prompt = "1MEVVIVTGMSGAGK"
13input_ids = torch.tensor(tokenizer.encode(prompt).ids).to(model.device)
14
15# forward pass
16logits = model(input_ids).logits
17
18# print output probabilities
19next_token_logits = logits[-1, :]
20next_token_probs = F.softmax(next_token_logits, dim=-1)
21for i in range(tokenizer.get_vocab_size(with_added_tokens=False)):
22 print(f"{tokenizer.id_to_token(i)}: {100 * next_token_probs[i].item():.2f} %")