Generation working. Todo: model edits; add model checkpoints. Also, we are only loading an editable model upon switching algs but we should load it when the page loads

algs/serac.py

@@ -306,13 +306,15 @@ class SERAC(EditableModel):
 
     def generate(self, *args, **kwargs):
         # input_text = self.replacement_tok.batch_decode(kwargs["input_ids"], skip_special_tokens=True)
+        if "max_new_tokens" not in kwargs:
+            kwargs["max_new_tokens"] = 20
         base_generate_fn = (
             self.model.forward if type(self.model) == BertClassifier 
-            else lambda *args, **kwargs: self.model.generate(*args, **kwargs, max_new_tokens=20)
+            else lambda *args, **kwargs: self.model.generate(*args, **kwargs)
         )
         cntr_generate_fn = (
             self.replacement.forward if type(self.replacement) == BertClassifier
-            else lambda *args, **kwargs: self.replacement.generate(*args, **kwargs, max_new_tokens=20)
+            else lambda *args, **kwargs: self.replacement.generate(*args, **kwargs)
         )
 
         # assert len(args) == 0, "Should only pass named arguments to generate()"

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import pandas as pd
 import time
 import importlib
+from torch.cuda import is_available as use_cuda
 
 import algs
 import config
@@ -17,6 +18,11 @@ EDIT_ALGS = [
     "LU: Lookup Cache",
 ]
 
+def generate(ids):
+    output_ids = st.session_state.editable_model.generate(input_ids=ids, max_new_tokens=20, min_length=1,
+                                                          num_return_sequences=1, num_beams=3)
+    return st.session_state.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
+
 def reset():
     st.session_state.edits.drop(st.session_state.edits.index, inplace=True)
     st.session_state.model_outputs.drop(st.session_state.edits.index, inplace=True)
@@ -28,10 +34,10 @@ def reset():
         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())
-        cfg = getattr(config, f"{alg_abbrv.lower()}_config")
+        st.session_state.config = getattr(config, f"{alg_abbrv.lower()}_config")
         st.session_state.editable_model = alg_class(
             st.session_state.model,
-            cfg,
+            st.session_state.config,
             lambda: copy.deepcopy(st.session_state.model),
         ).eval()
 
@@ -42,7 +48,10 @@ def apply_edit():
 
 def sample_model():
     input_str = str(test_input)
-    model_output = "blah blah blah"  ############## Actually sample the model
+    with st.spinner('Generating completion...'):
+        encoding = st.session_state.tokenizer(input_str, return_tensors="pt")
+        ids = encoding["input_ids"].to(st.session_state.device)
+        model_output = generate(ids)
     n_edits = len(st.session_state.edits)
     alg_name = st.session_state.alg_selector
     alg_abbrv = alg_name[:alg_name.index(":")]
@@ -55,9 +64,11 @@ if "init" not in st.session_state:
     st.session_state.edits = pd.DataFrame([], columns=["Edit input", "Edit label"])
     st.session_state.model_outputs = pd.DataFrame([], columns=["Input", "Output", "N edits", "Alg"])
     st.session_state.init = True
+    st.session_state.config = None
+    st.session_state.device = "cuda" if use_cuda() else "cpu"
     with st.spinner('Loading model...'):
         st.session_state.tokenizer = AutoTokenizer.from_pretrained("google/t5-large-ssm-nq")
-        st.session_state.model = AutoModelForSeq2SeqLM.from_pretrained("google/t5-large-ssm-nq").eval()
+        st.session_state.model = AutoModelForSeq2SeqLM.from_pretrained("google/t5-large-ssm-nq").to(st.session_state.device).eval()
         st.session_state.editable_model = None
 
 ########################