Getting Uncertainty results by using modified model output

[nehiridil]

Hey! I want to get uncertainty scores using model output as a parameter. We cannot make any modifications on the generated output from LLM as far as I understand. Is there a way to do that?

I would like to get the results from LLM, process it and than run uncertainty quantification.

Thanks!

[IINemo]

Hi, thank you for your question.
I have prepared example for you in the dev branch:
https://github.com/IINemo/lm-polygraph/blob/dev/examples/causal_lm_with_uncertainty.ipynb

This example has minimal changes to the code of the inference pipeline:

# LLM inference with uncertainty estimation

from transformers import AutoModelForCausalLM, AutoTokenizer

# ============== Addition Imports ===============
from lm_polygraph.estimators import MeanTokenEntropy
from lm_polygraph.stat_calculators import InferCausalLMCalculator, EntropyCalculator
from lm_polygraph.utils.causal_lm_with_uncertainty import CausalLMWithUncertainty
# ===============================================


device = "cuda"

# Loading standard LLM
llm = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
tokenizer.pad_token = tokenizer.eos_token
llm = llm.to(device)

# ======= Wrapping LLM with uncertainty estimator =========
stat_calculators = [InferCausalLMCalculator(tokenize=False),
                    EntropyCalculator()]
estimator = MeanTokenEntropy()
llm_with_uncertainty = CausalLMWithUncertainty(llm, tokenizer, stat_calculators, estimator)
# =========================================================

# Example prompt 
prompt = "Write a short story about a robot learning to paint.\n"

chat = [{"role": "user", "content": prompt}]
prompts = tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=False)
inputs = tokenizer(prompt, return_tensors="pt").to(device)

output = llm_with_uncertainty.generate(
    inputs,  
    max_new_tokens=200,
    temperature=0.7,
    do_sample=True
)

print("LLM output:")
print(tokenizer.decode(output.sequences[0], skip_special_tokens=True))

# ================ Printing uncertainty score ================
print("Uncertainty score: ", output.uncertainty_score)
# =============================================================

[IINemo]

Another low-level example could be found here:
https://github.com/IINemo/lm-polygraph/blob/main/examples/low_level_example.ipynb

Note that some UQ methods require run LLM multiple times, or they require some additional output from LLMs, like attention weights.

If you could share more info about your case, I could try further help you.

[nehiridil]

Hi!
Thank you so much for preparing this example code snippet. I am not sure if it solves my problem. I would like to explain my use-case a bit.
I am prompting LLMs for a classification task and then ask for a justification of their decision. I am introducing a json format as an output of LLM to extract the label and justification correctly. At the end I have this structure;

{
"label": 0,
"justification": "Explain the reason here"
}

I would like to run UQ on only justification, that's why I need to process the output from llm (extract only justification).

In addition to that,

I am experimenting basic UQ methods now using Whitebox models, the ones does not require any additional output like attention weights.

I have one more question on experimenting with different UQ methods at the same time. Currently I have function that iterates the UQ methods as you can see below;

from lm_polygraph.utils.manager import estimate_uncertainty

def estimate_uncertainties(model, ue_methods, input_text):
    scores = []
    for method_name in ue_methods.keys():
        ue = estimate_uncertainty(model, ue_methods[method_name], input_text=prompt)
        scores.append( 
            {
            'method': method_name,
            'uncertainty': ue.uncertainty
            }
        )
        torch.cuda.empty_cache()
    return scores, ue.generation_text

WIll the generation_text be the same for all of the iterations? I checked it manually but I am not sure if it is the case. Running the code takes some time and I would like to generate the text once and run UQ methods on generated text.

Thank you so much for your time!

[IINemo]

1. "I would like to run UQ on only justification" I think the simplest thing would be run classification and justification in two steps. While it is possible to get uncertainty only for justification via claim-level scores, you will need some manual code for representing tokens of justification as a claim. Moreover, I am not sure whether it will work well in terms of quality.
2. For the answer to your second question you can see the example here (Colab notebook): https://drive.google.com/file/d/1_TkY3Ak7gLNmivMuh3QW7BdZgOB-k2HG/view?usp=sharing