Views
No views yet
1from transformers import GPT2Tokenizer, GPT2LMHeadModel
2import torch
3
4# start and end tokens for generation
5START_TKN = "<|startoftext|>"
6END_TKN = "<|endoftext|>"
7
8# fine tuned on onion dataset w/ distilgpt2
9tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
10model = GPT2LMHeadModel.from_pretrained("distilgpt2")
11
12# use gpu if available
13device = "cpu"
14 if torch.cuda.is_available():
15 device = "cuda"
16
17model = model.to(device)
18
19# get 70th epoch (decent results)
20epoch = 70
21modelpath = f'distilgpt2_onion_{epoch}.pt'
22
23# load model
24model.load_state_dict(torch.load(modelpath))