import os import pickle as pkl from pathlib import Path from threading import Thread from typing import List, Tuple, Iterator from queue import Queue import spaces import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer MAX_MAX_NEW_TOKENS = 2048 DEFAULT_MAX_NEW_TOKENS = 1024 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.

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

""" EXAMPLES = [ ["What is the capital of France?", "You are a helpful assistant.", []], ["Explain the theory of relativity in simple terms.", "You are an expert physicist explaining concepts to a layman.", []], ["Write a short poem about artificial intelligence.", "You are a creative poet with a interest in technology.", []] ] 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", load_in_8bit=True) tokenizer = AutoTokenizer.from_pretrained(model_id) tokenizer.use_default_system_prompt = False # load the probe data # TODO compare accuracy and SE probe in different tabs/sections 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'] else: DESCRIPTION += "\n

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

" class CustomStreamer(TextIteratorStreamer): """ Streamer to also store hidden states in a queue. TODO check this works """ def __init__(self, tokenizer, skip_prompt: bool = False, skip_special_tokens: bool = False, **decode_kwargs): super().__init__(tokenizer, skip_prompt, skip_special_tokens, **decode_kwargs) self.hidden_states_queue = Queue() def put(self, value): if isinstance(value, dict) and 'hidden_states' in value: self.hidden_states_queue.put(value['hidden_states']) super().put(value) # Streamer claude # def generate( # message: str, # system_prompt: str, # chat_history: List[Tuple[str, 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, # ) -> Iterator[Tuple[str, str]]: # conversation = [] # if system_prompt: # conversation.append({"role": "system", "content": system_prompt}) # for user, assistant in chat_history: # conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}]) # 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) # streamer = CustomStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) # generation_kwargs = dict( # input_ids=input_ids, # max_new_tokens=max_new_tokens, # do_sample=True, # top_p=top_p, # top_k=top_k, # temperature=temperature, # repetition_penalty=repetition_penalty, # streamer=streamer, # output_hidden_states=True, # return_dict_in_generate=True, # ) # thread = Thread(target=model.generate, kwargs=generation_kwargs) # thread.start() # se_highlighted_text = "" # acc_highlighted_text = "" # for new_text in streamer: # hidden_states = streamer.hidden_states_queue.get() # # Semantic Uncertainty Probe # se_token_embeddings = torch.stack([layer[0, -1, :].cpu() for layer in hidden_states]) # se_concat_layers = se_token_embeddings.numpy()[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 # # Accuracy Probe # acc_token_embeddings = torch.stack([layer[0, -1, :].cpu() for layer in hidden_states]) # acc_concat_layers = acc_token_embeddings.numpy()[acc_layer_range[0]:acc_layer_range[1]].reshape(-1) # acc_probe_pred = acc_probe.predict_proba(acc_concat_layers.reshape(1, -1))[0][1] * 2 - 1 # se_new_highlighted_text = highlight_text(new_text, se_probe_pred) # acc_new_highlighted_text = highlight_text(new_text, acc_probe_pred) # se_highlighted_text += se_new_highlighted_text # acc_highlighted_text += acc_new_highlighted_text # yield se_highlighted_text, acc_highlighted_text @spaces.GPU def generate( message: str, chat_history: List[Tuple[str, 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, ) -> Iterator[str]: conversation = [] if system_prompt: conversation.append({"role": "system", "content": system_prompt}) for user, assistant in chat_history: conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}]) 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) streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True) generation_kwargs = dict( input_ids=input_ids, max_new_tokens=max_new_tokens, do_sample=True, top_p=top_p, top_k=top_k, temperature=temperature, repetition_penalty=repetition_penalty, streamer=streamer, output_hidden_states=True, return_dict_in_generate=True, ) # Generate without threading with torch.no_grad(): outputs = model.generate(**generation_kwargs) generated_tokens = outputs.sequences[0, input_ids.shape[1]:] generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True) # hidden states hidden = outputs.hidden_states # list of tensors, one for each token, then (batch size, sequence length, hidden size) # TODO do this loop on the fly instead of waiting for the whole generation se_highlighted_text = "" acc_highlighted_text = "" for i in range(1, len(hidden)): # Semantic Uncertainty Probe token_embeddings = torch.stack([generated_token[0, 0, :].cpu() for generated_token in hidden[i]]).numpy() # (num_layers, hidden_size) 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 # Accuracy Probe # acc_token_embeddings = torch.stack([layer[0, -1, :].cpu() for layer in hidden_states]) acc_concat_layers = token_embeddings[acc_layer_range[0]:acc_layer_range[1]].reshape(-1) acc_probe_pred = (1 - acc_probe.predict_proba(acc_concat_layers.reshape(1, -1))[0][1]) * 2 - 1 output_id = outputs.sequences[0, input_ids.shape[1]+i] output_word = tokenizer.decode(output_id) print(output_id, output_word, se_probe_pred, acc_probe_pred) se_new_highlighted_text = highlight_text(output_word, se_probe_pred) acc_new_highlighted_text = highlight_text(output_word, acc_probe_pred) se_highlighted_text += f" {se_new_highlighted_text}" acc_highlighted_text += f" {acc_new_highlighted_text}" yield se_highlighted_text, acc_highlighted_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 Dual 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.1, maximum=4.0, step=0.1, value=0.6) 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") # add spacing between probes and titles for each output with gr.Row(): with gr.Column(): title = gr.HTML("

Semantic Uncertainty Probe

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

Accuracy Probe

") acc_output = gr.HTML(label="Accuracy Probe") chat_history = gr.State([]) gr.Examples( examples=EXAMPLES, inputs=[message, system_prompt, chat_history], outputs=[se_output, acc_output], fn=generate, ) generate_btn.click( generate, inputs=[message, system_prompt, chat_history, max_new_tokens, temperature, top_p, top_k, repetition_penalty], outputs=[se_output, acc_output] ) # chat_interface = gr.ChatInterface( # fn=generate, # additional_inputs=[ # gr.Textbox(label="System prompt", lines=6), # gr.Slider( # label="Max new tokens", # minimum=1, # maximum=MAX_MAX_NEW_TOKENS, # step=1, # value=DEFAULT_MAX_NEW_TOKENS, # ), # gr.Slider( # label="Temperature", # minimum=0.1, # maximum=4.0, # step=0.1, # value=0.6, # ), # gr.Slider( # label="Top-p (nucleus sampling)", # minimum=0.05, # maximum=1.0, # step=0.05, # value=0.9, # ), # gr.Slider( # label="Top-k", # minimum=1, # maximum=1000, # step=1, # value=50, # ), # gr.Slider( # label="Repetition penalty", # minimum=1.0, # maximum=2.0, # step=0.05, # value=1.2, # ), # ], # stop_btn=None, # examples=[ # ["What is the capital of France?"], # ["Who landed on the moon?"], # ["Who is Yarin Gal?"] # ], # title="Llama-2 7B Chat with Streamable Semantic Uncertainty Probe", # description=DESCRIPTION, # ) # if __name__ == "__main__": # chat_interface.launch() if __name__ == "__main__": demo.launch()