Forced decoding with Huggingface Transformers

forced_decoding.py

from transformers import PrefixConstrainedLogitsProcessor

def create_processor_fn(ref_tokens_by_segment):
    def inner(batch_id, _):
        return ref_tokens_by_segment[batch_id]
    return inner

# ...

with tokenizer.as_target_tokenizer():
    tgt_encoded = tokenizer(tgt_lines)

logit_processor = PrefixConstrainedLogitsProcessor(create_processor_fn(tgt_encoded["input_ids"]), num_beams=5)

output = model.generate(**inputs, num_beams=5, logits_processor=[logit_processor], return_dict_in_generate=True, output_scores=True)
print(output.sequences_scores)