ericanthonymitchell/model-editing
0
1import streamlit as st2import pandas as pd3import time4import copy5import importlib6from torch.cuda import is_available as use_cuda7 8import algs9import config10from transformers import AutoTokenizer, AutoModelForSeq2SeqLM11import utils12 13 14EDIT_ALGS = [15 "MEND: Model editor networks using gradient decomposition",16 "SERAC: Semi-parametric editing with a retrieval-augmented counterfactual model",17 "ENN: Editable neural networks",18 "KE: KnowledgeEditor",19 "FT: Fine-tuning",20 "LU: Lookup Cache",21]22 23def get_alg_class(alg_abbrv):24 alg_module = importlib.import_module(f"algs.{alg_abbrv.lower()}")25 alg_class = getattr(alg_module, alg_abbrv.upper())26 return alg_class27 28def load_editable_model(alg_abbrv):29 alg_module = importlib.import_module(f"algs.{alg_abbrv.lower()}")30 alg_class = getattr(alg_module, alg_abbrv.upper())31 st.session_state.config = getattr(config, f"{alg_abbrv.lower()}_config")32 with st.spinner('Loading model...'):33 st.session_state.editable_model = alg_class(34 st.session_state.model,35 st.session_state.config,36 lambda: copy.deepcopy(st.session_state.model),37 ).eval()38 if "archive" in st.session_state.config:39 archive, st.session_state.config.archive = utils.load_archive(str(st.session_state.config.archive))40 print(f"Loading archive from {st.session_state.config.archive}")41 st.session_state.editable_model.load_state_dict(archive["model"])42 43def generate(ids):44 output_ids = st.session_state.editable_model.generate(input_ids=ids, max_new_tokens=20, min_length=1,45 num_return_sequences=1, num_beams=3)46 return st.session_state.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]47 48def reset():49 st.session_state.edits.drop(st.session_state.edits.index, inplace=True)50 st.session_state.model_outputs.drop(st.session_state.edits.index, inplace=True)51 52 selected_alg = st.session_state.alg_selector53 alg_abbrv = selected_alg[:selected_alg.index(":")]54 load_editable_model(alg_abbrv)55 56def apply_edit():57 st.session_state.edits.loc[len(st.session_state.edits)] = [str(edit_input), str(edit_label)]58 59 with st.spinner("Editing model..."):60 input_ids = st.session_state.tokenizer(str(edit_input), return_tensors="pt")["input_ids"].to(st.session_state.device)61 label_ids = st.session_state.tokenizer(str(edit_label), return_tensors="pt")["input_ids"].to(st.session_state.device)62 edit_sample = {"input_ids": input_ids, "labels": label_ids}63 st.session_state.editable_model, _ = st.session_state.editable_model.edit(edit_sample, detach_history=True)64 65def sample_model():66 input_str = str(test_input)67 with st.spinner('Generating completion...'):68 encoding = st.session_state.tokenizer(input_str, return_tensors="pt")69 ids = encoding["input_ids"].to(st.session_state.device)70 model_output = generate(ids)71 n_edits = len(st.session_state.edits)72 alg_name = st.session_state.alg_selector73 alg_abbrv = alg_name[:alg_name.index(":")]74 st.session_state.model_outputs.loc[len(st.session_state.model_outputs)] = [input_str, model_output, n_edits, alg_abbrv]75 76################################77#### Backend initialization ####78################################79if "init" not in st.session_state:80 st.session_state.edits = pd.DataFrame([], columns=["Edit input", "Edit label"])81 st.session_state.model_outputs = pd.DataFrame([], columns=["Input", "Output", "N edits", "Alg"])82 st.session_state.init = True83 st.session_state.device = "cpu" # "cuda" if use_cuda() else "cpu"84 with st.spinner('Loading model...'):85 st.session_state.tokenizer = AutoTokenizer.from_pretrained("google/t5-large-ssm-nq")86 st.session_state.model = AutoModelForSeq2SeqLM.from_pretrained("google/t5-large-ssm-nq").to(st.session_state.device).eval()87 # There is a "Loading model..." spinner in load_editable_model88 alg_abbrv = "MEND" # Default initial alg of dropdown selector89 load_editable_model(alg_abbrv)90 91########################92#### Interface code ####93########################94 95st.title("Language Model Editing")96st.markdown("**Note: this HF space is currently under development and doesn't actually work yet!**")97st.markdown("The goal of this demo is to give you a sense of the *abilities* and *limitations* of existing methods for **editing** pre-trained language models. **Model editing** algorithms use a single input-output pair to update a pre-trained model's behavior for that input (and ideally, related inputs).")98st.markdown("This demo uses a [T5-large](https://huggingface.co/google/t5-large-ssm-nq) model fine-tuned on [Natural Questions](https://arxiv.org/pdf/2002.08910.pdf) as the base pre-trained model.")99st.write("You can choose from a variety of algorithms for model editing in the dropdown below. At the bottom of the page, you can query the model for whatever input you want before/after editing.")100st.markdown("***")101 102col1, col2 = st.columns([5,1])103with col1:104 alg_selector = st.selectbox("Editing algorithm:", EDIT_ALGS, key="alg_selector", on_change=reset)105with col2:106 st.text("ㅤ")107 st.button("Clear edits", on_click=reset)108 109st.write("Edits applied so far:")110st.table(st.session_state.edits)111 112col1, col2, col3 = st.columns([3, 2, 1])113with col1:114 edit_input = st.text_input("Edit input:", placeholder="e.g., 'What is the tallest mountain on Earth?'")115with col2:116 edit_label = st.text_input("Edit target:", placeholder="e.g., 'Denali'")117with col3:118 st.text("ㅤ")119 edit_button = st.button("Apply edit", on_click=apply_edit)120 121st.markdown("***")122 123if len(st.session_state.edits) == 0:124 title = "Input to sample from *unedited* model:"125else:126 title = f"Input to sample from *edited* model:"127col1, col2 = st.columns([5, 1])128with col1:129 test_input = st.text_input(title, placeholder="e.g., 'What is the earth's tallest mountain?'")130with col2:131 st.text("ㅤ")132 generate_button = st.button("Generate", on_click=sample_model)133 134st.write("Model generation history:")135st.table(st.session_state.model_outputs)