Constrain generation to a pre-defined vocabulary

🤗Transformers
1

Hi, I’m trying to use a multilingual BART to generate output in a target language, which I have extracted a vocab file specific to it. I want to constrain the possible outputs of the generate() function to entries in the vocab file only. Currently I see the processor NoBadWordsLogitsProcessor which allows to prevent certain entries from being selected.

I tried setting the bad words list for every other token not in the vocab file but somehow it didn’t work

mbart = AutoModelForSeq2SeqLM.from_pretrained(mbart_id)
toker = AutoTokenizer.from_pretrained(mbart_id)
vocab_file = '/path/to/vocab/file"
bad_words_ids = set(range(toker.vocab_size)) - set(torch.load(vocab_file).values())

generation_arguments = dict(
    length_penalty=1.0,
    max_length=256,
    no_repeat_ngram_size=3,
    early_stopping=True,
    num_beams=1,
    return_dict_in_generate=True,
    bad_words_ids=[[x] for x in (bad_words_ids)],
    forced_bos_token_id=toker.lang_code_to_id["En_XX"],
)
out_tmp = mbart.generate(**inputs,**generation_arguments)['sequences']
print(toker.batch_decode(out_tmp,skip_special_tokens=True,
                                clean_up_tokenization_spaces=False)[0]) # output looks gibberish, full of repetition (2-grams)

Does anyone have any advice ?

Edit: I tried preventing tokens in the opposite language only and not every remaining tokens, output is still not quite right but seems more legitimate now !