Add logic for loading models

app.py

@@ -1,7 +1,12 @@
 import streamlit as st
 import pandas as pd
 import time
+import importlib
+
 import algs
+import config
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
 
 EDIT_ALGS = [
     "MEND: Model editor networks using gradient decomposition",
@@ -9,9 +14,13 @@ EDIT_ALGS = [
     "ENN: Editable neural networks",
     "KE: KnowledgeEditor",
     "FT: Fine-tuning",
-    "LU: Lookup Cache"
+    "LU: Lookup Cache",
 ]
 
+tokenizer = None
+model = None
+editable_model = None
+
 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)
@@ -19,7 +28,28 @@ def reset():
     selected_alg = st.session_state.alg_selector
     selected_alg_idx = EDIT_ALGS.index(selected_alg)
 
-    ############# Need to reset the model here (and maybe show progress spinner?)
+    ############# TODO: show progress spinner
+    global tokenizer
+    global model
+    global editable_model
+
+    if tokenizer is None:
+        tokenizer = AutoTokenizer.from_pretrained("google/t5-large-ssm-nq")
+    if model is None:
+        model = AutoModelForSeq2SeqLM.from_pretrained("google/t5-large-ssm-nq").eval()
+    del editable_model
+
+    alg_name = st.session_state.alg_selector
+    alg_abbrv = alg_name[:alg_name.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")
+    editable_model = alg_class(
+        model,
+        cfg,
+        lambda: copy.deepcopy(model),
+    ).eval()
+
 
 def apply_edit():
     st.session_state.edits.loc[len(st.session_state.edits)] = [str(edit_input), str(edit_label)]

config.py

@@ -0,0 +1,130 @@
+from omegaconf import OmegaConf
+from torch.cuda import is_available as use_cuda
+
+model_config = {
+  "name": "google/t5-large-ssm-nq",
+  "class_name": "AutoModelForSeq2SeqLM",
+  "tokenizer_class": "AutoTokenizer",
+  "tokenizer_name": "google/t5-large-ssm-nq",
+  "inner_params": [
+    "encoder.block.22.layer.1.DenseReluDense.wi.weight",
+    "encoder.block.22.layer.1.DenseReluDense.wo.weight",
+    "encoder.block.23.layer.1.DenseReluDense.wi.weight",
+    "encoder.block.23.layer.1.DenseReluDense.wo.weight",
+    "decoder.block.22.layer.2.DenseReluDense.wi.weight",
+    "decoder.block.22.layer.2.DenseReluDense.wo.weight",
+    "decoder.block.23.layer.2.DenseReluDense.wi.weight",
+    "decoder.block.23.layer.2.DenseReluDense.wo.weight",
+  ],
+  "pt": None,
+  "small_name": "t5-small",
+}
+
+ft_config = OmegaConf.create({
+  "device": "cuda" if use_cuda() else "cpu",
+  "edit_lr": 5e-6,
+  "train_base": False,
+  "ft": {
+    "verbose": False,
+    "max_edit_steps": 100,
+    "time_limit": None,
+    "locality": {
+      "enabled": False,
+      "oracle": True,
+      "cedit": 1e-2,
+      "batch_size": 1,
+    },
+    "rank": None,
+    "opt": "RMSprop",
+    "init_std": 0.01,
+  },
+  "model": model_config,
+})
+
+lu_config = OmegaConf.create({
+  "device": "cuda" if use_cuda() else "cpu",
+  "lu": {
+    "threshold": 2.75,
+    "onehot_logit": 1,
+  },
+  "model": model_config,
+})
+
+ke_config = OmegaConf.create({
+  "device": "cuda" if use_cuda() else "cpu",
+  "train_base": False,
+  "lr": 1e-5,
+  "model": model_config,
+})
+
+enn_config = OmegaConf.create({
+  "device": "cuda" if use_cuda() else "cpu",
+  "lr": 1e-5,
+  "edit_lr": 1e-2,
+  "lr_lr": 1e-3,
+  "train_base": True,
+  "grad_clip": 100,
+  "dropout": 0,
+  "no_grad_layers": None,
+  "enn": {
+    "first_order": False,
+    "n_edit_steps": 1,
+  },
+  "model": model_config,
+})
+
+mend_config = OmegaConf.create({
+  "device": "cuda" if use_cuda() else "cpu",
+  "lr": 1e-6,
+  "edit_lr": 1e-4,
+  "lr_lr": 1e-4,
+  "train_base": True,
+  "grad_clip": 100,
+  "dropout": 0,
+  "no_grad_layers": None,
+  "gtn": {
+    "one_sided": False,
+    "n_hidden": 1,
+    "hidden_dim": None,
+    "init": "id",
+    "norm": True,
+    "combine": True,
+    "x_only": False,
+    "delta_only": False,
+    "act": "relu",
+    "rank": 1920,
+    "mlp_class": "IDMLP",
+    "shared": True,
+    "descent": False,
+  },
+  "model": model_config,
+})
+
+serac_config = OmegaConf.create({
+  "device": "cuda" if use_cuda() else "cpu",
+  "lr": 1e-5,
+  "edit_lr": 1e-2,
+  "lr_lr": 0,
+  "train_base": False,
+  "grad_clip": 100,
+  "dropout": 0,
+  "no_grad_layers": None,
+  "rep": {
+    "cls_name": "distilbert-base-cased",
+    "cls_class": "AutoModel",
+    "supervised": "true",
+    "cos": False,
+    "freeze": None,
+    "square": True,
+    "bound_embeds": False,
+    "use_all_negatives": False,
+    "freeze_cntr": False,
+    "dist_heads": 1,
+    "cross_attend": False,
+    "lora": None,
+    "soft_weighting": False,
+    "checkpoint_grad": False,
+    "cache_embeds": True,
+  },
+  "model": model_config,
+})

requirements.txt

@@ -1,5 +1,6 @@
 allennlp
 git+https://github.com/eric-mitchell/higher@master # For in-place functional models
+omegaconf
 pandas
 streamlit
 torch