Views
No views yet
1from transformers import BartForConditionalGeneration, PreTrainedTokenizerFast, AutoTokenizer
2
3model = BartForConditionalGeneration.from_pretrained("KBLab/bart-base-swedish-cased")
4tok = AutoTokenizer.from_pretrained("KBLab/bart-base-swedish-cased")
5
6model.eval()
7
8input_ids = tok.encode(
9 "Jag har ätit en utsökt <mask> på restaurang vid <mask> .", return_tensors="pt"
10)
11
12# Simple greedy search
13output_ids = model.generate(
14 input_ids,
15 min_length=15,
16 max_length=25,
17 num_beams=1,
18 do_sample=False,
19)
20tok.decode(output_ids[0])
21# '</s><s> Jag har ätit en utsökt middag på restaurang vid havet på restaurang vid havet på restaurang vid havet.</s>'
22
23
24# Sampling
25output_ids = model.generate(
26 input_ids,
27 min_length=15,
28 max_length=20,
29 num_beams=1,
30 do_sample=True,
31)
32tok.decode(output_ids[0])
33#'</s><s> Jag har ätit en utsökt god mat som de tagit in på restaurang vid avröjda</s>'
34
35
36# Beam search
37output_ids = model.generate(
38 input_ids,
39 min_length=15,
40 max_length=25,
41 no_repeat_ngram_size=3,
42 num_beams=8,
43 early_stopping=True,
44 do_sample=True,
45 num_return_sequences=6
46)
47tok.decode(output_ids[0])
48# '</s><s> Jag har ätit en utsökt middag på restaurang vid havet. Jag har varit ute och gått en sväng.</s><pad><pad>'
49
50
51# Diverse beam generation
52output_ids = model.generate(
53 input_ids,
54 min_length=50,
55 max_length=100,
56 no_repeat_ngram_size=3,
57 num_beams=8,
58 early_stopping=True,
59 do_sample=False,
60 num_return_sequences=6,
61 num_beam_groups=8,
62 diversity_penalty=2.0,
63)
64tok.decode(output_ids[0])
65# '</s><s> Jag har ätit en utsökt middag på restaurang vid havet på restaurang. Jag har varit på restaurang i två dagar... Jag..,..!!!.. Så.. Nu.. Hej.. Vi.. Här.</s>'
66