{ "cells": [ { "cell_type": "markdown", "id": "9b42d4f5", "metadata": { "id": "9b42d4f5" }, "source": [ "# Uncertainty Analysis and Calibration for LLMs\n", "\n", "Author: Mehmet Emre Tiryaki (JHU)\n", "\n", "This lab notebook is organized around three parts:\n", "\n", "- **Part I — Uncertainty quantification:** Methods that measure the uncertainty in the answer of a language model.\n", "\n", "- **Part II — Evaluating uncertainty quantification:** Metrics that score uncertainty quantification methods. This is a tricky concept, the metrics here are concerned with how well uncertainty quantification scores align with accuracy.\n", "\n", "- **Part III — Calibration:** Methods that rescale an uncertainty quantification score to match better with the accuracy.\n", "\n", "**Your job: complete every cell marked `### TODO`.**" ] }, { "cell_type": "markdown", "id": "36727b81", "metadata": { "id": "36727b81" }, "source": [ "## 1. Setup\n", "\n", "Please don't forget to connect to a runtime with a sufficient GPU, T4 works here." ] }, { "cell_type": "code", "execution_count": null, "id": "522a286d", "metadata": { "id": "522a286d", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "022a39b8-71ab-44f1-9805-dd845d98ff7f" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ " Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m11.2/11.2 MB\u001b[0m \u001b[31m71.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m555.1/555.1 kB\u001b[0m \u001b[31m20.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m48.9/48.9 MB\u001b[0m \u001b[31m18.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25h Building wheel for rouge_score (setup.py) ... \u001b[?25l\u001b[?25hdone\n", "GPU 0: Tesla T4 (UUID: GPU-7d944ab8-80e1-6371-6aa1-7a198c3a3b7e)\n" ] } ], "source": [ "!pip install -q -U transformers datasets accelerate rouge_score\n", "!nvidia-smi -L" ] }, { "cell_type": "markdown", "id": "53c9325a", "metadata": { "id": "53c9325a" }, "source": [ "We are using Qwen2.5-3B-Instruct for this lab as the language model we are working on. The next two cells will download the model and setup some helpful functions." ] }, { "cell_type": "code", "execution_count": null, "id": "fbbdc581", "metadata": { "id": "fbbdc581", "colab": { "base_uri": "https://localhost:8080/", "height": 493, "referenced_widgets": [ "7103e1ad65f54e12b92fc98f66d38098", "21984568641a46f78b795aa203a4d027", "23b4877569314d228839ad7db4a210d1", "0ac7ff31066b4f7894bb1224e9c2eb58", "11dfda2a41484b49be9cc902a99e51b0", "e7711195948b44f4bb75c818e2c86e82", "8811325234124f06b11125a5c5388ca5", "32b146c325494d049e0ddf7b0a7a1732", "9434ac52bd094673a22f4de1a2d86466", "8544d282d6334bedbbd2af1b20d41968", "a0357b921cd84d57bc5a856cf87ce169", "e717bfd0d2cf404797b5abce24107edf", "a4b9269ee34a40bbbb231ff140884535", "a18ab8bc0b24411e9bd1cb9528b9cf2c", "215bacdc9e904705ac10e5b978d1c9e2", "89db3cbacbbd4f46812ea2525a122f37", "05a725cbd0354b8a9cb0d362e8a4f8f1", "103533dca50b43cc895723204b531a34", "14a265ebf1aa432bae0d3e22512aa290", "5d5496f645be4f46a7cde50956014f9c", "266472b6c631463ca065f38c5e5b6425", "b6a6a55bc5274469b123bb2cd7ca0f61", "e72441549fa048b1aa6563536b65084f", "ba286075bb574c0abb6fac90f7524e3a", "4d0f30dd6ba245828226124a56553499", "4f011326f28542aa94ed919c64cfcc9c", "9a7f44c1067f4dbab4d515c8f3b03d4d", "a0ef66a2fecf4a6cb21c0f99f1f43bff", "71a056894ce3438d9b9dc320366198a5", "6b59d32734904f558272e6a6f5bacc9a", "01aa973ba5cd460a84095119c488e64c", "8d39e3ac14c6416d85c7939784b8535d", "8a1c7e76b18e4492bfda508cecf8a54e", "9eea72f6e8da4ddaa8b9a35a25519121", "2678832320184c849f714e5d9bcc05a4", "189426f646884052be3abc10b45cd245", "ff4cd8ea16864b22a5dcec1c11009072", "07c629492e6e404bba00ee71c0fb370f", "565468c0cebc4e589b5aa7da2a868322", "3bc794ca96914d7db8855685fef2de63", "789caf0616d249cea779187fe59b82a2", "63a6c53619f449a1953007775424e4ae", "f08270ad274b4a68a56b9060f542281f", "f88d070070124d959f2b5d8d577f340a", "3583bae78c2a4b21bfc4ff7e62b33731", "9f62ee46f104435ca2455f26c67eb244", "ee629ac2441a4f9d963cde0f51f8e236", "55be8381eea04b268c2c01c1f6322610", "61c8a67f897245e58264c633d81e9102", "a10543b9f469427686d81b1fe4d4152e", "90295d65b56642a3894bc7b427fb1606", "0bd205955b9b4267b7ae2e9b257aadd7", "1af07954d6ab4c6cac19cb74af65aa2b", "f353bedbdf7041cf830e2bb5e09e1b2e", "5cb7ecf8706d4b12bd70cd8e09ffa2c0", "69918082f4454cf9bbb993238de1412c", "33e70069e46246ba947e4568419d4552", "ab0c1f08109140b2b6367af712998d0c", "539d43da6c1f4b8eb2426e82ae5ead41", "c94fa3beb7394ce39c394b2ef802e745", "967809c8825d4bfa97693015be89b352", "955aa8c4190f4388af853ebaeaf00887", "4c2a8c5d005f44a2abe28874d6f71850", "e70b3a60f0de4ad39e6ea962a195b792", "6265c246e7b24e83b1b37a9f9dea5cfe", "75794d0f48de44a7b40e152b6c1ee676", "3b3d0765591b407798b1d55d888da81b", "d79a39892a4249a2ae8ae0a79fe75ab7", "9d8b73ec452f444cb93de6fa33a069f4", "6ed8601b504c40db9ef25b1f86be990c", "d73a204f38ce40818ec5a1b64d0eb3db", "81a7e34772524b9d88633e37059b5453", "2017f7c4db74416d86aa2aae16e8d96f", "97aac36b9a464f56b2488ea4083771d1", "261e6dad2b504b768e6f158137026569", "20cc24dba7334746bfa85d5a5ed7d8f3", "9d086bc281fb415e964a2abea61a0f9b", "5b8150031a7844949e9cddc847d8b951", "d7729eedfc2449f8b4f5fe52df5dbe90", "655d8262d4db4a0a812c42efd0b94fa1", "2ab9fa1b1b954ffdaef68f6e19e0aef3", "78d8dd41009e44859b4f352fbecde212", "5b80643e2ab4409c93fc146a20ea466a", "504b992d77a543a8a819e16e10e990ad", "9d872b63905b465c839ad7a2de2d3fff", "316ddeaedb9d43589b4a139f7455085f", "eec83e83e63a4b24805c5324b7f66290", "423614d43fee41d69ace440d2c64a061", "92c364f520fa4638a584aed1e862f7ff", "c7d45069039a4570849256478fd373f1", "cef2dfdc4a4b49b2b6103fcfdeb5d46f", "7e0870ba336d4760b976841f298a3972", "791172b9e0d84d5380ff96eb22f39192", "e59e97516bee46619f294e1ccf711647", "4976d340ee7b46ff81dfbfe40b097507", "58d9cd9a5e284bed95ab79ab3625bd7e", "f84a28c95553465a8daf8b8f61098e36", "5487e7108cf1468eab9465b2ec3ee538", "7623653f31e749ad8af61e61fcf37231", "48fe1b70d8c4420dab104fdc3e78ed15", "bef0676d9392449598de1d3f61517c5f", "61c5de07b4874b0389a156980a0f2254", "7be7b6f77cf0400aa8712e39d7646167", "8e25d08318b148c9ad6d6614e20177e2", "13f9a9dc4e3047f5aedb6800d6748881", "983956e797bf41ae8833356763d8ef8c", "2fb2243089194bea887f43e39267754e", "acb6ed4b32954253be26404705777b81", "e16325a5f7534380939af0a008c19def", "2ee465f8d3d74f2f9cfaaab9eba293e0" ] }, "outputId": "a0663656-5fb6-44ef-9624-333ce381404a" }, "outputs": [ { "output_type": "stream", "name": "stderr", "text": [ "/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:134: UserWarning: \n", "Error while fetching `HF_TOKEN` secret value from your vault: 'Requesting secret HF_TOKEN timed out. Secrets can only be fetched when running from the Colab UI.'.\n", "You are not authenticated with the Hugging Face Hub in this notebook.\n", "If the error persists, please let us know by opening an issue on GitHub (https://github.com/huggingface/huggingface_hub/issues/new).\n", " warnings.warn(\n", "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", "WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n" ] }, { "output_type": "display_data", "data": { "text/plain": [ "config.json: 0%| | 0.00/661 [00:00 answer examples\n", " + [{\"role\": \"user\", \"content\": f\"Question: {question}\"}] # the question to answer\n", " )\n", " return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n", "\n", "def generate(question, n=1, sample=False, temperature=1.0, max_new_tokens=64, fewshot_messages=None):\n", " # (list of n answer strings, list of their per-token log-probabilities)\n", " fewshot_messages = FEWSHOT_MESSAGES if fewshot_messages is None else fewshot_messages\n", " enc = tokenizer(chat_prompt(question, fewshot_messages), return_tensors=\"pt\").to(device)\n", " prompt_len = enc.input_ids.shape[1]\n", "\n", " gen_kwargs = dict(max_new_tokens=max_new_tokens, do_sample=sample,\n", " num_return_sequences=n, pad_token_id=tokenizer.pad_token_id,\n", " return_dict_in_generate=True, output_scores=True)\n", " if sample:\n", " gen_kwargs.update(temperature=temperature, top_p=0.9, top_k=50)\n", "\n", " with torch.no_grad():\n", " out = model.generate(**enc, **gen_kwargs)\n", " transition = model.compute_transition_scores(out.sequences, out.scores, normalize_logits=True)\n", "\n", " new_tokens = out.sequences[:, prompt_len:]\n", " texts, logprobs = [], []\n", " for i in range(new_tokens.shape[0]):\n", " row = new_tokens[i]\n", " eos = (row == tokenizer.eos_token_id).nonzero()\n", " cut = int(eos[0]) if len(eos) > 0 else len(row)\n", " texts.append(tokenizer.decode(row[:cut], skip_special_tokens=True).strip())\n", " logprobs.append(transition[i, :cut].float().cpu().tolist())\n", " return texts, logprobs\n", "\n", "def hidden_states(question, answer, fewshot_messages=None):\n", " # array (num_layers, hidden_size): the model's internal state for this question+answer\n", " fewshot_messages = FEWSHOT_MESSAGES if fewshot_messages is None else fewshot_messages\n", " enc = tokenizer(chat_prompt(question, fewshot_messages) + answer, return_tensors=\"pt\").to(device)\n", " with torch.no_grad():\n", " out = model(**enc, output_hidden_states=True)\n", " idx = max(enc.input_ids.shape[1] - 2, 0)\n", " return torch.stack([layer[0, idx, :] for layer in out.hidden_states], dim=0).float().cpu().numpy()" ] }, { "cell_type": "markdown", "id": "34d653d1", "metadata": { "id": "34d653d1" }, "source": [ "## 2. Data and records\n", "\n", "TriviaQA is short-answer trivia dataset, we will use this dataset for the rest of the lab. We will use N_TEST number of questions for our test split, N_CALIB number of questions to calibrate the uncertainty quantification methods. For each questions we generate N_SAMPLES number of answers.We also sample a few **train** examples to build the\n", "few-shot prompt that will be used later.\n", "\n", "This cell will generate the answers using the numbers and the model that was determined before. We store these along with their log-probs in a dict for ease of use, you can inspect what is in the dict in the build_qa function." ] }, { "cell_type": "code", "execution_count": null, "id": "28e479e9", "metadata": { "id": "28e479e9", "colab": { "base_uri": "https://localhost:8080/", "height": 305, "referenced_widgets": [ "bac66ec0e6e5453289b927ecdd30d12d", "b6305c93facd44809570b9da6644b2f1", "8695753c6d0b4b9fbf96edb831ae7388", "4afb14ec08e44f4ba19204c201eaa647", "6bfce0addd294a33885827b4e0cedc76", "f921d43ba3934f199947ce350f3f75ed", "5fd292d7d0f0442f9bcff60c47529f61", "c6891c5b2645446280296cdf09136e32", "a069419c36f1492b9126c3424615f49d", "c7f61eabcf8c4c6f9ca66cabee4fee1a", "83ffdf5ef2ce4f329b56d77cf934d250", "96c0e319a4094389ab6ea9f9f01b21a5", "7a14af69bf6c4437884a41c8d0404bf4", "b85bec22a9dc40b9a7f3627d36b67918", "9a1b3e327d614df19f04809f3adcedef", "22795686dd0c4978bc879a54f00f7b67", "55e0cfeb099f4358aed4aa1b23e21181", "ca193079c550486881a7b877a57d848b", "c552e825599a4e139bed87283678347d", "9f15d5b0f1e64696b762bb206feefe5a", "61ded5db3c654a16aca9b4c37a720c9a", "890c4c78daa54ca09ce5cb3840d48f1f", "660bfbe65abc4e3293e3073f0fd5e012", "041ffe5b1c194812a866e7ca33829968", "285109f878cc40e2b6a2ef10ec8ca762", "f2fa7dad370d4c4ba5054356c9f37fb7", "d257a779c8464237928fefe6955c1f82", "4e79a8a32c6043459e1d6f40883a8d2f", "e74a587e146a4188a5f9fcc5ff516773", "395623293e494066916bc3214132a45c", "a0aaf2d650a34bea9ea40429f2235bce", "ed915e3478774cd18a677759367c9f3b", "464056ed5d2c40499d9c58e18b02baae", "6e24d27cb3344516841c427548d89e45", "745797fdc90346d582b1104f9bdcac56", "7be357788c3e4a2a8c2080a9dd68a33d", "39d67dd7c1c14216b060ebc0d52ce356", "b7cb21582cc7481a87fcc83d7698b598", "506b3cc149f34aa68ae1885894331f82", "adc9e8f4353e419d8395ae26c1c0f3d8", "604e8a676a354959b1dfe586aad3f4ce", "48c5d80c294948aab27261404ddd1815", "cf140e73df0e42d69440aeba121defdf", "1d11024172684a068fd942a59bfa146f", "a88661d894dc49869b7055210ccde1b5", "134bbf749fd548da84b7782b5254211a", "56dbf26c45664b339f8b39dcaf30495d", "aa14acdc35af4694ba37fe09add66679", "c816c396a8c84dc8a97d0950a3a217c1", "c0eeea01ef71462fa6af0ddc0dc37e12", "748d19a894754751915384376adc793d", "32dac79cf785435a9efcf2f89195c2ef", "530c913caa5f46a5a0e0e18d4a0a31c4", "d45615c4146c40e8b93de01b7c96bc0c", "a685429b8b5e46b7850c771730049a10", "5731929b60cb49389e3952f30b9eefe7", "ac10d34379aa4f80816e210222940c28", "9c07ec5a19104033917c3bb7e42fbe88", "33b9bce3665f458c98656cd814a1a675", "d0f3eb1619b14195b8bf46f0230abe35", "a9e3b0ae4ec94520a174516fe681a112", "a6964342599941ac89ea5bc04a96b290", "5ab1687b90c24bd8b30a525b0a31c4f8", "48714856a8044de3a42298039e4b8e0f", "b1443e37817d403e8f90f18f6009f230", "4cf69c42b6d24d98881a68bfe1f02649", "ef73b5b206174e0c90838f18341f7451", "29c4381d143b4c0ca413bc1423699f24", "72e587d08920467bb209ba33a27c9255", "4d1fe185e86e484e834d3244faa03266", "0358746ff7314db58e7229ad88ff8ca1", "1464a5bfcc8d4175b5d29e86fe88bdbf", "3b77905339ce4f318ea4089d778bc5cb", "f9c5d4818bc141a0bd3ea915df6a8c2b", "b0eb0c416992426f86352e6500106075", "174fbdf5f5014d5b9366ea73982ebd21", "9d1b2d3cd90545d88ad94cc9a2ff94d5", "24e8f933e09849c4a295b1123e639104", "410459ac62c24191b2fd0ece5da7e13a", "81d8166701e9476abfc02c3ebaee7b64", "0e1bbe9ab51e4692a2227e0a50afc7e8", "50abb09d696c478caca27f1fb2b9776a", "fe745b6dad364c22af1feb9c681d7999", "d080a7f4009a479a99cd63db1539ce94", "65e68517d1ff4ab6a7736aaae7a0acf9", "a5873f9e2fa5406c8cbfebb0a7dae900", "9be3dd5549404edfa2c5f59c8a4c929f", "0b83b5fa2cbf4883b26a47cc71c94277", "b94c73a73d354d1fb7fef920b8258170", "fa36579bc04242a6907f7d573463cac5", "c152a78db1e44aa9b978bdd54822ff5f", "efa3bc15e29749dba08a183a30654618", "dccb5c5cc81148cf87b86a06fec5a169", "f5c56e4c022842fcb1e9e29e7b61c377", "5f3e5bf506914f3b96cd9d66ad4c8eb6", "2d6f7b335c3d4c1292fde72450a9ed9f", "4cad8b745c8b4d63b42b6710822e3906", "5a2beb4448814b69b1699710ec1e6716", "334d0c6611404c508458a206fa70dd32" ] }, "outputId": "186f7f3e-50bb-4e62-c726-b3fad2194678" }, "outputs": [ { "output_type": "display_data", "data": { "text/plain": [ "README.md: 0%| | 0.00/26.7k [00:00= N_TEST + N_CALIB:\n", " break\n", "\n", "calib_items = items[:N_CALIB]\n", "test_items = items[N_CALIB:]\n", "\n", "def build_qa(item, with_hidden=True):\n", " # a dict with the question, references, greedy answer, samples, and their log-probs\n", " question = item[\"question\"]\n", " greedy_text, greedy_logprobs = generate(question, n=1, sample=False)\n", " sample_text, sample_logprobs = generate(question, n=N_SAMPLES, sample=True,\n", " temperature=SAMPLING_TEMPERATURE)\n", " qa = {\n", " \"question\": question,\n", " \"answers\": item.get(\"answers\"),\n", " \"greedy_answer\": greedy_text[0],\n", " \"greedy_logprobs\": greedy_logprobs[0],\n", " \"samples\": sample_text,\n", " \"sample_logprobs\": sample_logprobs,\n", " }\n", " if with_hidden:\n", " qa[\"hidden_states\"] = hidden_states(question, greedy_text[0])\n", " return qa\n", "\n", "test_qas = [build_qa(it) for it in test_items]\n", "calib_qas = [build_qa(it) for it in calib_items]" ] }, { "cell_type": "code", "execution_count": null, "id": "Pqds-XVFtL2S", "metadata": { "id": "Pqds-XVFtL2S" }, "outputs": [], "source": [ "dataset[0]" ] }, { "cell_type": "markdown", "id": "a3d74a33", "metadata": { "id": "a3d74a33" }, "source": [ "Here are the first few test records: the question, how many *distinct* answers the model\n", "gave across the samples, and the samples themselves. Familiarize yourself with the dataset and the type of answers our LLM tends to give." ] }, { "cell_type": "code", "execution_count": null, "id": "fd0223ea", "metadata": { "id": "fd0223ea" }, "outputs": [], "source": [ "import pandas as pd\n", "pd.set_option(\"display.max_colwidth\", None) # show full text; Colab clips wide cells by default\n", "\n", "pd.DataFrame([{\n", " \"question\": qa[\"question\"],\n", " \"distinct samples\": len({s.lower() for s in qa[\"samples\"]}),\n", " \"samples\": qa[\"samples\"],\n", "} for qa in test_qas[:8]])" ] }, { "cell_type": "markdown", "id": "18546b20", "metadata": { "id": "18546b20" }, "source": [ "## 3. Correctness\n", "\n", "Every uncertainty number in this lab is judged against a correctness label: was the model's\n", "answer right or wrong? So before measuring uncertainty we need a way to decide correctness\n", "automatically, by comparing the model's answer to one or more reference answers.\n", "\n", "The obvious choice, exact string match, is too strict for free-form text. The model\n", "might say *\"JFK\"* when the reference is *\"John F. Kennedy\"*, *\"Paris\"* vs *\"Paris, France\"*,\n", "*\"in 1944\"* vs *\"1944\"* and so on. So we need a softer notion of \"same answer\"\n", "and there's a spectrum of them.\n", "\n", "**Lexical overlap/Fuzzy Matching**: compare the words or characters of the answer to the reference label. These are cheap but has a not so small chance of failing when it comes to paraphrasing and synonyms. Most of these came from language translation.\n", "\n", "- **Token-F1** *(your TODO, and what we use)*: Treat each answer as a bag of words. Precision is\n", " the share of the answer's words found in the reference, recall is the share of the reference's\n", " words the answer covers. F1 score is their harmonic mean.We call an answer\n", " correct when **Token-F1 against the best-matching reference > T** where you choose T. We will use > 0.5 for this lab. It\n", " rewards word overlap regardless of meaning, \"not guilty\"* vs *\"guilty\"* score higher than true\n", " synonyms like \"car\" and \"automobile\".\n", "- **Rouge-L** *(provided)*: based on the **longest common subsequence** — the longest run of words\n", " appearing in both answers in the same order. Same harmonic-mean idea\n", " as Token-F1, but order-sensitive, so it rewards getting the sequence right. Common in\n", " summarization.\n", "- **BLEU** *: an n-gram precision metric from machine translation.\n", " It checks how many 1-, 2-, 3-, and 4-word groups of the answer appear in the reference, takes a\n", " geometric average of those precisions, and multiplies by a brevity penalty so an answer\n", " can't win just by being short. Because it leans on longer n-grams, it's designed for full\n", " sentences.\n", "\n", "**Semantic / model-based**: compare meaning using a learned model. The methods may be more accurate then lexical similarity based methods but they have their own inacurracies as they are dependent on the quality of the model used and they are more expensive to run.\n", "\n", "- **BERTScore**: Instead of matching exact words, it embeds every\n", " token of the answer and the reference with a language model (e.g. finetuned BERT) and matches them by\n", " cosine similarity of their embeddings, then aggregates into a precision/recall/F1. Because\n", " it works in embedding space, *\"car\"* and *\"automobile\"* match even though they share no letters, it aims to capture semantic information directly.\n", "\n", "- Further along this axis: an **LLM-as-a-judge** (ask a strong model \"is this answer correct given\n", " the reference?\"), or **human annotators**, which is the most faithful but also the most expensive.\n", "\n", "We use **Token-F1 > 0.5** for the lab following some previous work as instant, but it's imperfect. Keep that in mind and note that which correctness metric you use is one of the knobs you can experiment with later." ] }, { "cell_type": "code", "execution_count": null, "id": "c8e5a576", "metadata": { "id": "c8e5a576" }, "outputs": [], "source": [ "def normalize(text):\n", " # the text lowercased, with articles and punctuation stripped (for fair matching)\n", " text = text.lower()\n", " text = re.sub(r\"\\b(a|an|the)\\b\", \" \", text)\n", " text = re.sub(r\"[^a-z0-9 ]\", \" \", text)\n", " return re.sub(r\"\\s+\", \" \", text).strip()" ] }, { "cell_type": "code", "execution_count": null, "id": "d1d60c53", "metadata": { "id": "d1d60c53" }, "outputs": [], "source": [ "def token_f1(prediction, reference):\n", " ### TODO: Token-F1 between prediction and reference.\n", "\n", " # Algorithm:\n", " # 1. Normalize both strings and split them into word lists.\n", " # 2. Count the words they share, counting a repeated word only as many times as it\n", " # appears in both.\n", " # 3. Precision is shared over the prediction length; recall is shared over the\n", " # reference length.\n", " # 4. Return the harmonic mean of precision and recall (0 if nothing is shared).\n", "\n", " raise NotImplementedError(\"Implement Token-F1\")" ] }, { "cell_type": "code", "execution_count": null, "id": "1b1bf981", "metadata": { "id": "1b1bf981" }, "outputs": [], "source": [ "from rouge_score import rouge_scorer\n", "_rouge_scorer = rouge_scorer.RougeScorer([\"rougeL\"], use_stemmer=False)\n", "\n", "def rouge_l(prediction, reference):\n", " # Rouge-L F-measure in [0, 1] (first argument is the reference, second the prediction)\n", " return _rouge_scorer.score(reference, prediction)[\"rougeL\"].fmeasure\n", "\n", "def is_correct(prediction, references, threshold=0.5):\n", " # True if the prediction matches any reference alias well enough (Token-F1 > 0.5 default)\n", " if isinstance(references, str):\n", " references = [references]\n", " return max((token_f1(prediction, r) for r in references), default=0.0) >= threshold" ] }, { "cell_type": "markdown", "id": "34aa41c4", "metadata": { "id": "34aa41c4" }, "source": [ "Let's check the accuracy of our model (using its greedy answers) on our test set using our definition of correctness and inspect some of them. How much was actually incorrect, how much is due to errors caused by Token-F1 as the correctness measure? You might also notice some bad reference answers." ] }, { "cell_type": "code", "execution_count": null, "id": "8d81696c", "metadata": { "id": "8d81696c" }, "outputs": [], "source": [ "import pandas as pd\n", "pd.set_option(\"display.max_colwidth\", None)\n", "\n", "test_correct = np.array([is_correct(qa[\"greedy_answer\"], qa[\"answers\"]) for qa in test_qas])\n", "calib_correct = np.array([is_correct(qa[\"greedy_answer\"], qa[\"answers\"]) for qa in calib_qas])\n", "print(f\"greedy accuracy on test: {test_correct.mean():.1%}\")\n", "\n", "df = pd.DataFrame([\n", " {\"question\": qa[\"question\"], \"model answer\": qa[\"greedy_answer\"],\n", " \"reference\": qa[\"answers\"][0], \"correct\": bool(ok)}\n", " for qa, ok in zip(test_qas, test_correct)])\n", "\n", "right, wrong = df[df.correct], df[~df.correct]\n", "print(\"Examples the model got RIGHT:\"); display(right.sample(min(3, len(right))))\n", "print(\"\\nExamples the model got WRONG:\"); display(wrong.sample(min(3, len(wrong))))\n", "\n", "acc = test_correct.mean()\n", "plt.figure(figsize=(5, 1.3))\n", "plt.barh([\"test\"], [acc], color=\"#4c9f70\")\n", "plt.barh([\"test\"], [1 - acc], left=[acc], color=\"#d96c6c\")\n", "plt.xlim(0, 1); plt.xticks([0, 0.5, 1]); plt.title(f\"greedy accuracy = {acc:.1%} (green = correct)\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "c8f7df14", "metadata": { "id": "c8f7df14" }, "source": [ "# Part I — Uncertainty Quantification\n", "\n", "An **uncertainty quantification (UQ) method** is a function that maps a model's answer\n", "to a number, the *uncertainty score*, where larger means the model is more uncertain. The methods differ in what they read to compute the\n", "number and their costs.\n", "\n", "We may classify the methods into black-box (output text only), grey-box\n", "(token probabilities), and white-box (hidden states) depending on what the method needs as an input. Then, we have the following methods at hand in this notebook:\n", "\n", "- **Sequence-likelihood methods** *(grey-box)*:\n", " token confidence (perplexity), predictive entropy, length-normalized entropy\n", " (Malinin & Gales).\n", "\n", "- **Self-evaluation methods**: ask the model about its own answer: P(True) *(grey-box;\n", " Kadavath)*, verbalized confidence *(black-box; Lin et al., Tian et al.)*,\n", " LLM-as-a-judge *(grey-box)*.\n", "- **Semantic Similarity Methods**: sample several answers and measure agreement in some space: lexical similarity *(black-box; Kuhn, after Fomicheva et al.)*, semantic entropy *(grey-box or black-box; Kuhn)*.\n", "- **Others**: semantic entropy probe *(white-box; Kossen\n", " et al.)*." ] }, { "cell_type": "markdown", "id": "9f823358", "metadata": { "id": "9f823358" }, "source": [ "### I.I Sequence-likelihood methods\n", "\n", "These read the log-probabilities the model already assigns to tokens. Write the greedy\n", "answer as tokens $w_1,\\dots,w_T$, and a sampled answer $s_n$ with sequence\n", "log-probability $\\log p(s_n\\mid x) = \\sum_t \\log p(w^{(n)}_t \\mid w^{(n)}_{ array of per-question uncertainty scores (test set)\n", "\n", "def show_uncertainty(items, qas, bins=20):\n", " # items: list of (name, scores). Show each method's score distribution and the questions\n", " # it rates least / most uncertain.\n", " fig, axes = plt.subplots(1, len(items), figsize=(3.7 * len(items), 2.6))\n", " for ax, (name, scores) in zip(np.atleast_1d(axes), items):\n", " ax.hist(scores, bins=bins, color=\"#4c78a8\")\n", " ax.set_title(name, fontsize=9); ax.set_xlabel(\"uncertainty\"); ax.set_ylabel(\"count\")\n", " plt.tight_layout(); plt.show()\n", " for name, scores in items:\n", " o = np.argsort(scores)\n", " lo, hi = qas[o[0]], qas[o[-1]]\n", " print(name)\n", " print(f\" least uncertain -> {lo['greedy_answer']} (Q: {lo['question']})\")\n", " print(f\" most uncertain -> {hi['greedy_answer']} (Q: {hi['question']})\")\n", "\n", "# Compute the three sequence-likelihood scores on the test set and look at the values.\n", "uncertainty[\"token_confidence\"] = np.array([token_confidence(qa) for qa in test_qas])\n", "uncertainty[\"predictive_entropy\"] = np.array([predictive_entropy(qa) for qa in test_qas])\n", "uncertainty[\"length_norm_entropy\"] = np.array([length_norm_entropy(qa) for qa in test_qas])\n", "\n", "show_uncertainty([(n, uncertainty[n]) for n in\n", " [\"token_confidence\", \"predictive_entropy\", \"length_norm_entropy\"]], test_qas)" ] }, { "cell_type": "markdown", "id": "b4b2d620", "metadata": { "id": "b4b2d620" }, "source": [ "### I.II Self Evaluation Methods\n", "\n", "\n", "- **Verbalized confidence** a very naive method where we generate the answer to the question \"How confident are you?\" from the LLM from 1 to 100.\n", "\n", " $c \\in [0,100]$ and use $U = 1 - c/100$.\n", "\n", "- **LLM-as-a-judge**: another naive method, ask \"is this correct?\" and read\n", " the probability of \"Yes\" without supplying sample answers or few-shot examples.\n", "\n", "- **P(True)** *(your exercise, below)*: using an instruction tuned model show the model the question, some candidate answers by sampling them, the proposed answer to check confidence, and ask if the proposed answer is `A) True` / `B) False`, and read $p(\\texttt{A})$.\n", " To ensure correct formatting, thus intended probabilities for the A and B tokens we provide few-shot turns built from sampled **train** questions, each labelled `A`/`B` by\n", " correctness.\n", " $$U = 1 - p(\\texttt{A}\\mid \\text{prompt}).$$" ] }, { "cell_type": "code", "execution_count": null, "id": "833f1b3b", "metadata": { "id": "833f1b3b" }, "outputs": [], "source": [ "def next_token_prob(prompt, words):\n", " # probability that the model's next token is one of `words`\n", " enc = tokenizer(prompt, return_tensors=\"pt\").to(device)\n", " with torch.no_grad():\n", " logits = model(**enc).logits[0, -1, :]\n", " probs = torch.softmax(logits.float(), dim=-1)\n", " total = 0.0\n", " for word in words:\n", " ids = tokenizer.encode(word, add_special_tokens=False)\n", " if ids:\n", " total += probs[ids[0]].item()\n", " return total\n", "\n", "def verbalized_conf(qa):\n", " # the model's self-reported confidence in [0, 1]\n", " user = (f\"Question: {qa['question']}\\n\"\n", " f\"Your answer: {qa['greedy_answer']}\\n\"\n", " f\"Output only your confidence as an integer 0-100.\")\n", " prompt = tokenizer.apply_chat_template(\n", " [{\"role\": \"system\", \"content\": \"Output only an integer from 0 to 100.\"},\n", " {\"role\": \"user\", \"content\": user}],\n", " tokenize=False, add_generation_prompt=True)\n", " enc = tokenizer(prompt, return_tensors=\"pt\").to(device)\n", " with torch.no_grad():\n", " out = model.generate(**enc, max_new_tokens=8, do_sample=False,\n", " pad_token_id=tokenizer.pad_token_id)\n", " text = tokenizer.decode(out[0, enc.input_ids.shape[1]:], skip_special_tokens=True)\n", " match = re.search(r\"\\d+\", text)\n", " if not match:\n", " return 0.5\n", " return max(0.0, min(100.0, float(match.group()))) / 100.0\n", "\n", "def judge_conf(qa):\n", " # p(Yes) among {Yes, No}: the model's judgement that its own answer is correct\n", " user = (f\"Question: {qa['question']}\\n\"\n", " f\"Answer: {qa['greedy_answer']}\\n\"\n", " f\"Is the answer correct? Reply Yes or No.\")\n", " prompt = tokenizer.apply_chat_template(\n", " [{\"role\": \"system\", \"content\": \"You grade answers. Reply Yes or No.\"},\n", " {\"role\": \"user\", \"content\": user}],\n", " tokenize=False, add_generation_prompt=True)\n", " p_yes = next_token_prob(prompt, [\" Yes\", \"Yes\", \"yes\"])\n", " p_no = next_token_prob(prompt, [\" No\", \"No\", \"no\"])\n", " if p_yes + p_no == 0:\n", " return 0.5\n", " return p_yes / (p_yes + p_no)" ] }, { "cell_type": "code", "execution_count": null, "id": "3751c89a", "metadata": { "id": "3751c89a" }, "outputs": [], "source": [ "# Few-shot turns for P(True), built from sampled train questions and labelled A/B by\n", "# correctness. This pins the A/B output format so p(A) is meaningful.\n", "P_TRUE_NUM_FEWSHOT = 5\n", "\n", "def ptrue_user_block(question, brainstormed, proposed):\n", " # Format one P(True) prompt as text: the question, the brainstormed candidate answers,\n", " # the proposed answer to judge, and the \"A) True / B) False\" choice. This is the user\n", " # message you send when asking the model whether `proposed` is correct.\n", " lines = [f\"Question: {question}\", \"Brainstormed Answers:\"]\n", " lines += [b.strip() for b in brainstormed]\n", " lines += [f\"Possible answer: {proposed}\",\n", " \"Is the possible answer:\", \"A) True\", \"B) False\",\n", " \"Respond with just A or B.\"]\n", " return \"\\n\".join(lines)\n", "\n", "def build_ptrue_messages(indices):\n", " # Build the few-shot examples for P(True): for each train question, generate answers,\n", " # wrap them in a ptrue_user_block, and label it \"A\" (true) or \"B\" (false) by actual\n", " # correctness. Returns chat turns to prepend so the model replies with a single A/B token.\n", " msgs = []\n", " for fi in indices:\n", " ex = train_data[fi]\n", " question = ex[\"question\"]\n", " references = ex[\"answer\"][\"aliases\"] or [ex[\"answer\"][\"value\"]]\n", " most_likely = generate(question, n=1, sample=False)[0][0]\n", " samples = generate(question, n=N_SAMPLES, sample=True, temperature=SAMPLING_TEMPERATURE)[0]\n", " block = ptrue_user_block(question, [most_likely] + samples, most_likely)\n", " msgs.append({\"role\": \"user\", \"content\": block})\n", " msgs.append({\"role\": \"assistant\", \"content\": \"A\" if is_correct(most_likely, references) else \"B\"})\n", " return msgs\n", "\n", "ptrue_fewshot_idx = random.Random(1).sample(range(len(train_data)), P_TRUE_NUM_FEWSHOT)\n", "PTRUE_MESSAGES = build_ptrue_messages(ptrue_fewshot_idx) # few-shot turns, reused by p_true_conf below" ] }, { "cell_type": "code", "execution_count": null, "id": "e03945ad", "metadata": { "id": "e03945ad" }, "outputs": [], "source": [ "def p_true_conf(qa):\n", " ### TODO: P(True) confidence: p(\"A\") normalized against p(\"B\").\n", " # Use the helpers defined just above this cell:\n", " # - ptrue_user_block(question, candidates, proposed) formats the A/B prompt block.\n", " # Use the samples plus the greedy answer as the candidates, and the greedy answer\n", " # as the proposed answer to judge.\n", " # - PTRUE_MESSAGES are the few-shot turns that pin the A/B format. Put them before\n", " # your block, then apply the chat template (as in build_ptrue_messages / verbalized_conf).\n", " # - next_token_prob(prompt, words) reads the next-token probability of those words\n", " # (pass both spaced and unspaced variants, like judge_conf does for Yes/No).\n", " # Read p(\"A\") and p(\"B\"), then return p(A) / (p(A) + p(B)).\n", " raise NotImplementedError(\"Implement P(True)\")" ] }, { "cell_type": "markdown", "id": "6a317b9b", "metadata": { "id": "6a317b9b" }, "source": [ "This cell checks some answers and what these methods gave as uncertainty scores." ] }, { "cell_type": "code", "execution_count": null, "id": "43b6d86f", "metadata": { "id": "43b6d86f" }, "outputs": [], "source": [ "import pandas as pd\n", "\n", "pd.DataFrame([{\n", " \"answer\": qa[\"greedy_answer\"][:24],\n", " \"correct\": bool(is_correct(qa[\"greedy_answer\"], qa[\"answers\"])),\n", " \"verbal\": round(verbalized_conf(qa), 2),\n", " \"judge\": round(judge_conf(qa), 2),\n", " \"P(True)\": round(p_true_conf(qa), 2),\n", "} for qa in test_qas[:8]])" ] }, { "cell_type": "markdown", "id": "62ccb1fa", "metadata": { "id": "62ccb1fa" }, "source": [ "### I.III Semantic Similarity Methods\n", "\n", "Sample several answers, measure some sort of agreement.\n", "\n", "- **Lexical similarity**: agreement in *surface form*, the average pairwise Rouge-L\n", " among the $N$ samples:\n", " $$U_{\\text{lex}} = 1 - \\frac{1}{\\binom{N}{2}}\\sum_{i