Constrain generation to a pre-defined vocabulary
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 workmbart = 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 !