import streamlit as st import torch from transformers import AutoTokenizer, AutoModelForCausalLM from transformers import DynamicCache USE_GPU = torch.cuda.is_available() @st.cache_resource def load_model(): import torch model_name = 'google/gemma-2-9b-it' dtype = torch.bfloat16 if USE_GPU else torch.float16 llm = { 'tokenizer': AutoTokenizer.from_pretrained(model_name), 'model': AutoModelForCausalLM.from_pretrained( model_name, device_map="auto" if USE_GPU else "cpu", torch_dtype=dtype, attn_implementation='eager' ) } llm['model'].eval() return llm def type_assistant_response(): if 'messages' not in st.session_state or st.button("Start a new conversation"): st.session_state['messages'] = [{"role": "user", "content": ""}] st.session_state['msg_in_progress'] = "" messages = st.session_state.messages def rewind_to(i): st.session_state.messages = st.session_state.messages[:i+1] st.session_state['msg_in_progress'] = st.session_state.messages[-1]['content'] for i, message in enumerate(st.session_state.messages[:-1]): with st.chat_message(message["role"]): st.markdown(message["content"]) st.button("Edit", on_click=rewind_to, args=(i,), key=f"rewind_to_{i}") # Display message-in-progress in chat message container last_role = messages[-1]["role"] with st.chat_message(last_role): label = "Your message" if last_role == "user" else "Assistant response" msg_in_progress = st.text_area(label, placeholder="Clicking the buttons below will update this field. You can also edit it directly; press Ctrl+Enter to apply changes.", height=300, key="msg_in_progress") if msg_in_progress is None: msg_in_progress = "" messages[-1]['content'] = msg_in_progress def append_token(word): messages[-1]['content'] = st.session_state['msg_in_progress'] = ( msg_in_progress + word ) allow_multi_word = st.checkbox("Allow multi-word predictions", value=False) response = continue_messages( messages=messages, n_branch_tokens=5, n_future_tokens=2 ) continuations = response['continuations'] for i, (col, continuation) in enumerate(zip(st.columns(len(continuations)), continuations)): token = continuation['doc_text'] with col: if not allow_multi_word and ' ' in token[1:]: token = token[0] + token[1:].split(' ', 1)[0] # if not allow_multi_word: # import re # split_result = re.split(r'(\s+)', token, maxsplit=1) # assert len(split_result) == 3 # before_ws, token, after_ws = split_result # print(repr(split_result)) # if before_ws != '': # token = before_ws token_display = show_token(token) st.button(token_display, on_click=append_token, args=(token,), key=i, use_container_width=True) def send_message(): other_role = "assistant" if last_role == "user" else "user" st.session_state['messages'].append({"role": other_role, "content": ""}) st.session_state['msg_in_progress'] = "" st.button("Send", on_click=send_message) def show_token(token: str, escape_markdown=True) -> str: token_display = token.replace('\n', '↵').replace('\t', '⇥') if escape_markdown: for c in "\\`*_{}[]()#+-.!": token_display = token_display.replace(c, "\\" + c) return token_display def continue_messages(messages, n_branch_tokens, n_future_tokens): messages = [{"role": m.role, "content": m.content} for m in messages] if len(messages) == 0: raise ValueError("At least one message must be provided.") llm = load_model() model = llm['model'] tokenizer = llm['tokenizer'] generated_docs = continue_messages_inner(model, tokenizer, messages, n_branch_tokens, n_future_tokens) return { 'continuations': [dict(doc_text=doc) for doc in generated_docs] } def get_lookahead_sequences(model, tokenizer, hypotheses, n_branch_tokens, device): """ For each of the n_branch_tokens next tokens, generate most-likely next tokens and append back on. """ assert len(hypotheses.shape) == 2 assert hypotheses.shape[0] == 1 n_tokens_so_far = hypotheses.shape[1] past_key_values = DynamicCache() with torch.no_grad(): model_outs_onestep = model(hypotheses, output_hidden_states=True, past_key_values=past_key_values) branch_tokens = model_outs_onestep.logits[0, -1].topk(n_branch_tokens).indices # split the cache into n_branch_tokens reps. We pretend we're doing a "Beam search"... past_key_values.reorder_cache(torch.zeros((n_branch_tokens,), dtype=torch.long, device=device)) # Now call the model again, passing the kv cache, so we can continue generating. # Each of the n_branch_tokens next tokens will be considered as one sequence in a "batch". next_tokens_as_batch = branch_tokens.unsqueeze(1) assert next_tokens_as_batch.shape == (n_branch_tokens, 1) position_id_for_final_token = n_tokens_so_far cache_position = torch.full((1,), position_id_for_final_token, dtype=int, device=device) with torch.no_grad(): model_outs = model( next_tokens_as_batch, past_key_values=past_key_values, output_hidden_states=True, use_cache=True, # the cache surprisingly doesn't know the position of the last token cache_position=cache_position ) # Grab the single most likely token from each of the n_branch_tokens sequences next_token_logits = model_outs.logits[:, -1] vocab_size = model.config.vocab_size assert next_token_logits.shape == (n_branch_tokens, vocab_size), f"{next_token_logits.shape=}, {n_branch_tokens=}, {vocab_size=}" most_likely_token_ids = next_token_logits.argmax(dim=-1) # Stick them at the end of the branch tokens. assert most_likely_token_ids.shape == (n_branch_tokens,) lookahead_sequences = torch.cat([ branch_tokens.unsqueeze(1), most_likely_token_ids.unsqueeze(1) ], dim=1) assert lookahead_sequences.shape == (n_branch_tokens, 2) return lookahead_sequences, next_token_logits def continue_messages_inner(model, tokenizer, messages, n_branch_tokens, n_future_tokens): # Note: we're ignoring n_future_tokens right now since the old implementation was buggy. device = model.device tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt", continue_final_message=True).to(model.device) print(tokenizer.batch_decode(tokenized_chat, skip_special_tokens=False)) lookahead_sequences, next_token_logits = get_lookahead_sequences( model, tokenizer, tokenized_chat, n_branch_tokens, device) generated_docs = tokenizer.batch_decode(lookahead_sequences, skip_special_tokens=True) return generated_docs type_assistant_response()