Package: mbr
mbr 🔥mbradds Sampling-based Minimum Bayes Risk decoding toHugging Face transformers. Originally proposed byEikema & Aziz (2022), this technique is a risk-minimizing algorithm for generating text with a language model.Pronounce:ember/ˈɛm.bɚ/InstallationpipinstallmbrRequirements:Python >= 3.9PyTorchUsageThe main components ofmbrare:mbr.MBRGenerationMixin: overrides a model'sgeneratemethod to add MBR decoding.mbr.MBRGenerationConfig: specifies the parameters of MBR decoding, e.g., the number of samples to generate and the metric to optimize.1. Load a Hugging Face transformers modelModels need to inherit fromMBRGenerationMixinfor MBR decoding to work. Here's two ways to achieve this, using the Llama model as an example:Variant A:fromtransformersimportLlamaForCausalLMfrommbrimportMBRGenerationMixinclassMBRLlamaForCausalLM(MBRGenerationMixin,LlamaForCausalLM):passThen, you can useMBRLlamaForCausalLMas you would useLlamaForCausalLM:model=MBRLlamaForCausalLM.from_pretrained(...)Variant B:frommbrimportMBRmodel=MBR(LlamaForCausalLM).from_pretrained(...)2. Configure MBR decodingCreate anMBRConfigobject to pass to the model'sgeneratemethod:frommbrimportMBRConfigmbr_config=MBRConfig(num_samples=10,metric="chrf",)3. Generate text as usualCall the model'sgeneratemethod directly, or use the Pipeline API. Make sure to pass thembr_config, as well as the model's tokenizer.fromtransformersimportpipelinegenerator=pipeline("text-generation",model=model,tokenizer=tokenizer)output=generator("Hello,",mbr_config=mbr_config,tokenizer=tokenizer)How MBR decoding worksThe following research papers, among many others, provide a description of Sampling-based Minimum Bayes Risk decoding:Sampling-Based Approximations to Minimum Bayes Risk Decoding for Neural Machine Translation(Eikema & Aziz, EMNLP 2022)Understanding the Properties of Minimum Bayes Risk Decoding in Neural Machine Translation(Müller & Sennrich, ACL-IJCNLP 2021)In practice, MBR decoding is most commonly implemented as follows (on the example of machine translation):Instead of searching for the single most probable output sequence (e.g., using beam search), generate a number of samples.Score each sample against the other samples using a metric (e.g., BLEU).Return the sample with the highest score. Intuitively, this can be seen as returning the median of all samples.The terminology around MBR decoding varies:Term used in this codebaseAlternative termssamplescandidates, hypothesesreferencespseudo-references, evidencemetric scoreexpected utility(negative) expected risk, errorDetailsConfiguring the samplingThe generation of the samples can be customized by passing ageneration_configto thegeneratemethod or to the pipeline call:fromtransformersimportGenerationConfiggeneration_config=GenerationConfig.from_pretrained("mymodel",do_sample=True,num_beams=1,epsilon_cutoff=0.02,)model.generate(...,generation_config=generation_config)Separate set of referencesBy default, the samples themselves are used a references (or a subset of the samples ifnum_referencesis smaller thannum_samples).You could also sample the reference set independently, using a custom generation config for the references:fromtransformersimportGenerationConfigreferences_config=GenerationConfig.from_pretrained("mymodel",do_sample=True,num_beams=1,top_p=0.9,)model.generate(...,references_config=references_config)Choosing a metricBy default,mbrintegrates metrics via theHugging Face Evaluatelibrary.A full list of metrics is foundhere. Some typical choices are:ChrF(Popović, 2015)COMET(Rei et al., 2020)BLEURT(Sellam et al., 2020)In the MBR config, you can either specify the metric's name (e.g.,"chrf","comet") or pass anevaluate.Metricobject directly.Since different metrics output differently structured dicts, you need to specify themetric_output_fieldthat should be used as the metric score.fromevaluateimportloadmetric=load('chrf')mbr_config=MBRGenerationConfig(metric=metric,metric_output_field="score",# the ChrF metric returns a dict with a "score" field...)Customizing the metric computationInternally,mbrwill call the metric'scomputemethod to calculate the metric score for each sample.By default,mbrwill callcomputeseparately for each sample–reference pair.
Since this requires manycomputecalls, it can make sense to optimize the metric computation. Different metrics will require different optimization strategies.
To override the default way of calling the metric, define aMetricRunnerclass and pass it to thegeneratemethod:frommbrimportMetricRunnerclassMyMetricRunner(MetricRunner):def__call__(self,input_ids:torch.LongTensor,sample_ids:Tuple[torch.LongTensor],reference_ids:Tuple[torch.LongTensor],)->torch.FloatTensor:...# TODO: implement your efficient metric computation heremodel.generate(...,metric_runner=MyMetricRunner())ForCOMET, an optimized implementation is already provided inCometMetricRunner:frommbr.metrics.cometimportCometMetricRunnermbr_config=MBRGenerationConfig(...,metric="comet",metric_output_field="mean_score",)metric_runner=CometMetricRunner(mbr_config,tokenizer)model.generate(...,metric_runner=metric_runner)OptimizationsMBR decoding is notoriously slow.mbrimplements some optimizations:Cached encoder outputs: For encoder-decoder models, the encoder outputs are computed only once and reused during sampling.Cached metric: The metric is computed only once for each unique sample–reference pair (since there will be duplicate samples and references).Optimized COMET metric: Inspired byAmrhein & Sennrich (2022), sequence embeddings are cached and reused for all pairwise comparisons.Example scriptsTheexperimentsdirectory contains the code for reproductions of experiments from the following papers:MBR for (low-resource) machine translation(Müller & Sennrich, 2021)MBR with neural metrics and epsilon sampling for machine translation(Freitag et al., 2023)MBR for summarization(Bertsch et al., 2023)Related projectshttps://github.com/roxot/mbr-nmt: Original implementation (demo)https://github.com/ZurichNLP/understanding-mbr: MBR with Sockeyehttps://github.com/ZurichNLP/mbr-sensitivityandhttps://github.com/Unbabel/COMET#minimum-bayes-risk-decoding: COMET metric for MBRhttps://github.com/rainavyas/mbr_gec: MBR for Grammatical Error CorrectionChangelogv0.2.0Breaking change:RenameMBRGenerationConfigtoMBRConfigBreaking change:MetricRunnernow returns aMetricOutputdict instead of the raw tensor of scores.Make the size of the metric cache configurable viaMBRConfig.metric_cache_sizeAllow that the number of references can be larger than the number of samples (if generated separately from the samples).RemoveGenerationConfigas parent class ofMBRConfig