IINemo / lm-polygraph

MIT License
111 stars 21 forks source link

May I know how to run an ensemble-based uncertainty estimation method? #170

Closed he159ok closed 5 months ago

he159ok commented 6 months ago

Hi, I have successfully run the information-based and density-based methods.

May I know how to run the ensemble-based uncertainty estimation method?

IINemo commented 6 months ago

@rvashurin

he159ok commented 6 months ago

Thanks for your reply.

I have found a rough solution for the ensemble-based uncertainty estimation method, which I share below.

step 1. add below in an example's .py file.

import lm_polygraph.utils as utils

ensemble_model = None
if args.use_ensemble_ue:
    ensemble_model = utils.model.create_ensemble(
        model_paths=[model_name_or_path],
        mc=True,
        seed=1,
        mc_seeds=[2],
        ensembling_mode="pe",
        device="cuda:0",
        dropout_rate=0.1,
    )

before initializing an UEManager() object.

step 2. in the manager.py -> UEManager() -> call(), add below code

            if self.ensemble_model is not None:
                batch_stats["ensemble_model"] = self.ensemble_model
                if "ensemble_generation_params" not in batch_stats.keys():
                    batch_stats["ensemble_generation_params"] = {}

before the row of batch_stats = self.calculate(batch_stats, self.stat_calculators, inp_texts).