Incomplete editing implementation

algs/lu.py

@@ -67,7 +67,7 @@ class LU(EditableModel):
 
         return output
 
-    def edit(self, batch, condition=None):
+    def edit(self, batch, condition=None, detach_history=False):
         edit_model = self.model.eval()
         if "bert" in self.config.model.name.lower():
             _, encoder_states = self.model(**batch, output_hidden_states=True)

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import pandas as pd
 import time
+import copy
 import importlib
 from torch.cuda import is_available as use_cuda
 
@@ -28,13 +29,11 @@ def reset():
     st.session_state.model_outputs.drop(st.session_state.edits.index, inplace=True)
 
     selected_alg = st.session_state.alg_selector
-    selected_alg_idx = EDIT_ALGS.index(selected_alg)
-
+    alg_abbrv = selected_alg[:selected_alg.index(":")]
+    alg_module = importlib.import_module(f"algs.{alg_abbrv.lower()}")
+    alg_class = getattr(alg_module, alg_abbrv.upper())
+    st.session_state.config = getattr(config, f"{alg_abbrv.lower()}_config")
     with st.spinner('Loading model...'):
-        alg_abbrv = selected_alg[:selected_alg.index(":")]
-        alg_module = importlib.import_module(f"algs.{alg_abbrv.lower()}")
-        alg_class = getattr(alg_module, alg_abbrv.upper())
-        st.session_state.config = getattr(config, f"{alg_abbrv.lower()}_config")
         st.session_state.editable_model = alg_class(
             st.session_state.model,
             st.session_state.config,
@@ -44,7 +43,11 @@ def reset():
 def apply_edit():
     st.session_state.edits.loc[len(st.session_state.edits)] = [str(edit_input), str(edit_label)]
 
-    ############# Actually do the edit to the model
+    with st.spinner("Editing model..."):
+        input_ids = st.session_state.tokenizer(str(edit_input), return_tensors="pt")["input_ids"].to(st.session_state.device)
+        label_ids = st.session_state.tokenizer(str(edit_label), return_tensors="pt")["input_ids"].to(st.session_state.device)
+        edit_sample = {"input_ids": input_ids, "labels": label_ids}
+        st.session_state.editable_model, _ = st.session_state.editable_model.edit(edit_sample, detach_history=True)
 
 def sample_model():
     input_str = str(test_input)

config.py

@@ -24,6 +24,7 @@ ft_config = OmegaConf.create({
   "device": "cuda" if use_cuda() else "cpu",
   "edit_lr": 5e-6,
   "train_base": False,
+  "grad_clip": 100,
   "ft": {
     "verbose": False,
     "max_edit_steps": 100,