import os import pickle as pkl from pathlib import Path from threading import Thread from typing import List, Tuple, Iterator, Optional, Generator from queue import Queue import spaces import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer # TODO this is not as fast as it could be using generate function with 1 token at a time # TODO log prob output scaling highlighting instead? # TODO make it look nicer # TODO better examples. # TODO streaming output (need custom generation function because of probes) # TODO add options to switch between models, SLT/TBG, layers? # TODO full semantic entropy calculation MAX_MAX_NEW_TOKENS = 1024 DEFAULT_MAX_NEW_TOKENS = 100 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096")) DESCRIPTION = """

Llama-2 7B Chat with Uncertainty Probes

This Space demonstrates the Llama-2-7b-chat model with a semantic uncertainty probe.

This demo is based on our paper: "Semantic Entropy Probes: Robust and Cheap Hallucination Detection in LLMs" by Jannik Kossen*, Jiatong Han*, Muhammed Razzak*, Lisa Schut, Shreshth Malik and Yarin Gal.

The highlighted text shows the model's uncertainty in real-time:

Please see our paper for more details. NOTE: This demo is a work in progress.

""" EXAMPLES = [ ["What is the capital of France?", ""], ["Who landed on the moon?", ""], ["Who is Yarin Gal?", ""], ["Explain the theory of relativity in simple terms.", ""], ] if torch.cuda.is_available(): model_id = "meta-llama/Llama-2-7b-chat-hf" # TODO load the full model not the 8bit one? model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_id) tokenizer.use_default_system_prompt = False # load the probe data with open("./model/20240625-131035_demo.pkl", "rb") as f: probe_data = pkl.load(f) # take the NQ open one probe_data = probe_data[-2] se_probe = probe_data['t_bmodel'] se_layer_range = probe_data['sep_layer_range'] acc_probe = probe_data['t_amodel'] acc_layer_range = probe_data['ap_layer_range'] print(f"Loaded probes with layer ranges: {se_layer_range}, {acc_layer_range}") else: DESCRIPTION += "\n

Running on CPU 🥶 This demo does not work on CPU.

" @spaces.GPU def generate( message: str, system_prompt: str, max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, temperature: float = 0.6, top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1.2, ) -> Generator[Tuple[str, str], None, None]: conversation = [] if system_prompt: conversation.append({"role": "system", "content": system_prompt}) conversation.append({"role": "user", "content": message}) input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt") if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH: input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:] gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.") input_ids = input_ids.to(model.device) generation_kwargs = dict( do_sample=True, top_p=top_p, top_k=top_k, temperature=temperature, repetition_penalty=repetition_penalty, output_hidden_states=True, return_dict_in_generate=True, ) sentence_start_idx = input_ids.shape[1] sentence_token_count = 0 finished = False with torch.no_grad(): # highlight and return the prompt outputs = model.generate(**generation_kwargs, input_ids=input_ids, max_new_tokens=1) prompt_tokens = outputs.sequences[0, :input_ids.shape[1]] prompt_text = tokenizer.decode(prompt_tokens, skip_special_tokens=True) print(prompt_tokens, prompt_text) # hidden states hidden = outputs.hidden_states # last token embeddings (note this is the same as the token before generation given this is the prompt) token_embeddings = torch.stack([generated_token[0, -1, :].cpu() for generated_token in hidden[0]]).numpy() se_concat_layers = token_embeddings[se_layer_range[0]:se_layer_range[1]].reshape(-1) se_probe_pred = se_probe.predict_proba(se_concat_layers.reshape(1, -1))[0][1] * 2 - 1 acc_concat_layers = token_embeddings[acc_layer_range[0]:acc_layer_range[1]].reshape(-1) acc_probe_pred = acc_probe.predict_proba(acc_concat_layers.reshape(1, -1))[0][0] * 2 - 1 # accuracy probe is inverted wrt uncertainty se_new_highlighted_text = highlight_text(prompt_text, se_probe_pred) acc_new_highlighted_text = highlight_text(prompt_text, acc_probe_pred) se_highlighted_text = f"{se_new_highlighted_text}
" acc_highlighted_text = f"{acc_new_highlighted_text}
" while not finished: outputs = model.generate(**generation_kwargs, input_ids=input_ids, max_new_tokens=1) # this should only be the one extra token (equivalent to -1) generated_tokens = outputs.sequences[0, input_ids.shape[1]:] print(f"generated_tokens {generated_tokens}" ) # add to the conversation input_ids = torch.cat([input_ids, generated_tokens.unsqueeze(0)], dim=-1) # stop at the end of a sequence if generated_tokens[-1] == tokenizer.eos_token_id or input_ids.shape[1] > max_new_tokens: print("Finished") finished = True if generated_text != "": # do final prediction on the last generated text (one before the eos token) print("Predicting probes") hidden = outputs.hidden_states # hidden states = (num generated tokens, num layers, batch size, num tokens, hidden size) # last token embeddings token_embeddings = torch.stack([generated_token[0, -2, :].cpu() for generated_token in hidden[-1]]).numpy() se_concat_layers = token_embeddings[se_layer_range[0]:se_layer_range[1]].reshape(-1) se_probe_pred = se_probe.predict_proba(se_concat_layers.reshape(1, -1))[0][1] * 2 - 1 acc_concat_layers = token_embeddings[acc_layer_range[0]:acc_layer_range[1]].reshape(-1) acc_probe_pred = acc_probe.predict_proba(acc_concat_layers.reshape(1, -1))[0][0] * 2 - 1 print(f"se_probe_pred {se_probe_pred}, acc_probe_pred {acc_probe_pred}") se_new_highlighted_text = highlight_text(generated_text, se_probe_pred) acc_new_highlighted_text = highlight_text(generated_text, acc_probe_pred) se_highlighted_text += f" {se_new_highlighted_text}" acc_highlighted_text += f" {acc_new_highlighted_text}" sentence_start_idx += sentence_token_count sentence_token_count = 0 # decode the full generated text generated_text = tokenizer.decode(outputs.sequences[0, sentence_start_idx:], skip_special_tokens=True) print(f"generated_text: {generated_text}") sentence_token_count += 1 # TODO this should be when a factoid is detected rather than just punctuation. Is the SLT token always basically a period for the probes? if generated_text.endswith(('.', '!', '?', ';', '."', '!"', '?"')): print("Predicting probes") hidden = outputs.hidden_states # hidden states = (num generated tokens, num layers, batch size, num tokens, hidden size) # last token embeddings token_embeddings = torch.stack([generated_token[0, -1, :].cpu() for generated_token in hidden[-1]]).numpy() se_concat_layers = token_embeddings[se_layer_range[0]:se_layer_range[1]].reshape(-1) se_probe_pred = se_probe.predict_proba(se_concat_layers.reshape(1, -1))[0][1] * 2 - 1 acc_concat_layers = token_embeddings[acc_layer_range[0]:acc_layer_range[1]].reshape(-1) acc_probe_pred = acc_probe.predict_proba(acc_concat_layers.reshape(1, -1))[0][0] * 2 - 1 print(f"se_probe_pred {se_probe_pred}, acc_probe_pred {acc_probe_pred}") se_new_highlighted_text = highlight_text(generated_text, se_probe_pred) acc_new_highlighted_text = highlight_text(generated_text, acc_probe_pred) se_highlighted_text += f" {se_new_highlighted_text}" acc_highlighted_text += f" {acc_new_highlighted_text}" sentence_start_idx += sentence_token_count sentence_token_count = 0 generated_text = "" # yield se_highlighted_text + generated_text, acc_highlighted_text + generated_text yield se_highlighted_text + generated_text #, acc_highlighted_text + generated_text def highlight_text(text: str, uncertainty_score: float) -> str: if uncertainty_score > 0: html_color = "#%02X%02X%02X" % ( 255, int(255 * (1 - uncertainty_score)), int(255 * (1 - uncertainty_score)), ) else: html_color = "#%02X%02X%02X" % ( int(255 * (1 + uncertainty_score)), 255, int(255 * (1 + uncertainty_score)), ) return '{}'.format( html_color, text ) with gr.Blocks(title="Llama-2 7B Chat with Semantic Uncertainty Probes", css="footer {visibility: hidden}") as demo: gr.HTML(DESCRIPTION) with gr.Row(): with gr.Column(): message = gr.Textbox(label="Message") system_prompt = gr.Textbox(label="System prompt", lines=2) with gr.Column(): max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS) temperature = gr.Slider(label="Temperature", minimum=0.01, maximum=2.0, step=0.1, value=0.01) top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9) top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50) repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2) with gr.Row(): generate_btn = gr.Button("Generate") stop_btn = gr.Button("Stop") # Add spacing between probes gr.HTML("

") # with gr.Row(): with gr.Column(): title = gr.HTML("

Semantic Uncertainty Probe

") se_output = gr.HTML(label="Semantic Uncertainty Probe") # with gr.Column(): # make a box # title = gr.HTML("

Semantic Uncertainty Probe

") # se_output = gr.HTML(label="Semantic Uncertainty Probe") # Add spacing between columns # gr.HTML("
") # with gr.Column(): # title = gr.HTML("

Accuracy Probe

") # acc_output = gr.HTML(label="Accuracy Probe") gr.Examples( examples=EXAMPLES, inputs=[message, system_prompt], # outputs=[se_output, acc_output], outputs=[se_output], fn=generate, ) generate_event = generate_btn.click( generate, inputs=[message, system_prompt, max_new_tokens, temperature, top_p, top_k, repetition_penalty], # outputs=[se_output, acc_output] outputs=[se_output] ) stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[generate_event]) if __name__ == "__main__": demo.launch()