KB-BART

一个在瑞典语语料库上进行训练的模型，包含150亿个标记（约80GB的文本）。该模型是使用 Fairseq 进行训练的，并转换为与Huggingface兼容的格式。

训练代码可以在 here 找到。

使用方法

from transformers import BartForConditionalGeneration, PreTrainedTokenizerFast, AutoTokenizer

model = BartForConditionalGeneration.from_pretrained("KBLab/bart-base-swedish-cased")
tok = AutoTokenizer.from_pretrained("KBLab/bart-base-swedish-cased")

model.eval()

input_ids = tok.encode(
    "Jag har ätit en utsökt <mask> på restaurang vid <mask> .", return_tensors="pt"
)

# Simple greedy search
output_ids = model.generate(
    input_ids,
    min_length=15,
    max_length=25,
    num_beams=1,
    do_sample=False,
)
tok.decode(output_ids[0])
# '</s><s> Jag har ätit en utsökt middag på restaurang vid havet på restaurang vid havet på restaurang vid havet.</s>'


# Sampling
output_ids = model.generate(
    input_ids,
    min_length=15,
    max_length=20,
    num_beams=1,
    do_sample=True,
)
tok.decode(output_ids[0])
#'</s><s> Jag har ätit en utsökt god mat som de tagit in på restaurang vid avröjda</s>'


# Beam search
output_ids = model.generate(
    input_ids,
    min_length=15,
    max_length=25,
    no_repeat_ngram_size=3,
    num_beams=8,
    early_stopping=True,
    do_sample=True,
    num_return_sequences=6
)
tok.decode(output_ids[0])
# '</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>'


# Diverse beam generation
output_ids = model.generate(
    input_ids,
    min_length=50,
    max_length=100,
    no_repeat_ngram_size=3,
    num_beams=8,
    early_stopping=True,
    do_sample=False,
    num_return_sequences=6,
    num_beam_groups=8,
    diversity_penalty=2.0,
)
tok.decode(output_ids[0])
# '</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>'

致谢

我们衷心感谢HPC RIVR联合会（ www.hpc-rivr.si ）和EuroHPC JU（ eurohpc-ju.europa.eu/ ）提供了在信息科学研究所的HPC系统Vega上提供计算资源的资助。