Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4
5device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
6tokenizer = AutoTokenizer.from_pretrained("tuanle/GPT2_Poet")
7
8model = AutoModelForCausalLM.from_pretrained("tuanle/GPT2_Poet").to(device)
9text = "hỏi rằng nàng"
10
11input_ids = tokenizer.encode(text, return_tensors='pt').to(device)
12min_length = 60
13max_length = 100
14
15sample_outputs = model.generate(input_ids,pad_token_id=tokenizer.eos_token_id,
16 do_sample=True,
17 max_length=max_length,
18 min_length=min_length,
19 # temperature = .8,
20 # top_k= 100,
21 top_p = 0.8,
22 num_beams= 10,
23 # early_stopping=True,
24 no_repeat_ngram_size= 2,
25 num_return_sequences= 3)
26
27for i, sample_output in enumerate(sample_outputs):
28 print(">> Generated text {}\n\n{}".format(i+1, tokenizer.decode(sample_output.tolist(), skip_special_tokens=True)))
29 print('\n---')