init

app.py

@@ -0,0 +1,241 @@
+# https://huggingface.co/spaces/Glaciohound/LM-Steer
+
+import torch
+import streamlit as st
+import random
+import numpy as np
+import pandas as pd
+from lm_steer.models.get_model import get_model
+
+
+@st.cache_resource(show_spinner="Loading model...")
+def st_get_model(model_name, low_resource_mode):
+    device = torch.device("cuda:0") if torch.cuda.is_available() \
+        else torch.device("cpu")
+    model, tokenizer = get_model(
+        model_name, "final_layer", "multiply",
+        4,
+        1000, 1e-3, 1e-2, low_resource_mode
+    )
+    model.to_device(device)
+    ckpt = torch.load(f"checkpoints/{model_name}.pt", map_location=device)
+    model.load_state_dict(ckpt[1])
+    return model, tokenizer
+
+
+def word_embedding_space_analysis(model, tokenizer, dim):
+    matrix = model.steer.projector1.data[dim].matmul(
+        model.steer.projector2.data[dim].transpose(0, 1))
+    S, V, D = torch.linalg.svd(matrix)
+    embeddings = model.steer.lm_head.weight
+
+    data = []
+    for _i in range(10):
+        left_tokens = embeddings.matmul(D[_i]).argsort()[-20:].flip(0)
+        right_tokens = embeddings.matmul(D[_i]).argsort()[:20]
+
+        def filter_words(side_tokens):
+            output = []
+            for t in side_tokens:
+                word = tokenizer.decode([t])
+                if not word[0].isalpha() and word[1:].isalpha():
+                    output.append(word[1:]+"-")
+            return output
+
+        data.append([
+            ", ".join(filter_words(side_tokens))
+            for side_tokens in [left_tokens, right_tokens]
+        ])
+    st.table(pd.DataFrame(
+        data,
+        columns=["One Direction", "Another Direction"],
+        index=[f"Dim {_i}" for _i in range(10)],
+    ))
+
+
+def main():
+    # set up the page
+    random.seed(0)
+    title = "LM-Steer: Word Embeddings Are Steers for Language Models"
+    st.set_page_config(
+        layout="wide",
+        page_title=title,
+        page_icon="🛞",
+    )
+    st.title(title)
+    '''
+    Live demo for the paper ["**LM-Steer: Word Embeddings Are Steers for
+    Language Models**"](https://arxiv.org/abs/2305.12798) (**ACL 2024
+    Outstanding Paper Award**) by Chi Han, Jialiang Xu, Manling Li, Yi Fung,
+    Chenkai Sun, Nan Jiang, Tarek Abdelzaher, Heng Ji. GitHub repository:
+    https://github.com/Glaciohound/LM-Steer.
+    '''
+    st.subheader("Overview")
+    st.image('https://raw.githubusercontent.com/Glaciohound/LM-Steer'
+             '/refs/heads/main/assets/overview_fig.jpg')
+    '''
+    Language models (LMs) automatically learn word embeddings during
+    pre-training on language corpora. Although word embeddings are usually
+    interpreted as feature vectors for individual words, their roles in
+    language model generation remain underexplored. In this work, we
+    theoretically and empirically revisit output word embeddings and find that
+    their linear transformations are equivalent to steering language model
+    generation styles. We name such steers LM-Steers and find them existing in
+    LMs of all sizes. It requires learning parameters equal to 0.2% of the
+    original LMs' size for steering each style.
+    '''
+
+    # set up the model
+    st.divider()
+    st.divider()
+    st.subheader("Select a model:")
+    '''
+    Due to resource limits, we are only able to provide a few models for
+    steering. You can also refer to the Github repository:
+    https://github.com/Glaciohound/LM-Steer for hosting larger models.
+    '''
+    col1, col2 = st.columns(2)
+    st.session_state.model_name = col1.selectbox(
+        "Select a model to steer",
+        [
+            "gpt2",
+            "gpt2-medium",
+            "gpt2-large",
+            "EleutherAI/pythia-70m",
+            "EleutherAI/pythia-160m",
+            "EleutherAI/pythia-410m",
+            # "EleutherAI/pythia-1b", "EleutherAI/pythia-1.4b",
+            # "EleutherAI/pythia-2.8b", "EleutherAI/pythia-6.9b",
+            # "EleutherAI/gpt-j-6B",
+        ],
+    )
+    low_resource_mode = True if st.session_state.model_name in (
+        "EleutherAI/pythia-1.4b", "EleutherAI/pythia-2.8b",
+        "EleutherAI/pythia-6.9b", "EleutherAI/gpt-j-6B",
+    ) else False
+    model, tokenizer = st_get_model(
+        st.session_state.model_name, low_resource_mode)
+    num_param = model.steer.projector1.data.shape[1] ** 2 / 1024 ** 2
+    total_param = sum(p.numel() for _, p in model.named_parameters()) / \
+        1024 ** 2
+    ratio = num_param / total_param
+    col2.write(f"Steered {num_param:.1f}M out of {total_param:.1f}M "
+               "parameters, ratio: {:.2%}".format(ratio))
+
+    # steering
+    steer_range = 4.
+    steer_interval = 0.5
+    st.subheader("Enter a sentence and steer the model")
+    st.session_state.prompt = st.text_input(
+        "Enter a prompt",
+        st.session_state.get("prompt", "My life")
+    )
+    # col1, col2, col3 = st.columns(3, gap="medium")
+    col1, col2, col3 = st.columns([2, 2, 1], gap="medium")
+    sentiment = col1.slider(
+        "Sentiment", -steer_range, steer_range, 3.0, steer_interval)
+    detoxification = col2.slider(
+        "Detoxification Strength", -steer_range, steer_range, 0.0,
+        steer_interval)
+    max_length = col3.number_input("Max length", 50, 300, 50, 50)
+    col1, col2, col3, _ = st.columns(4)
+    randomness = col2.checkbox("Random sampling", value=False)
+
+    if "output" not in st.session_state:
+        st.session_state.output = ""
+    if col1.button("Steer and generate!", type="primary"):
+        steer_values = [detoxification, 0, sentiment, 0]
+        st.session_state.output = model.generate(
+            st.session_state.prompt,
+            steer_values,
+            seed=None if randomness else 0,
+            min_length=0,
+            max_length=max_length,
+            do_sample=True,
+        )
+    analyzed_text = \
+        st.text_area("Generated text:", st.session_state.output, height=200)
+
+    # Analysing the sentence
+    st.divider()
+    st.divider()
+    st.subheader("Analyzing Styled Texts")
+    '''
+    LM-Steer also serves as a probe for analyzing the text. It can be used to
+    analyze the sentiment and detoxification of the text. Now, we proceed and
+    use LM-Steer to analyze the text in the box above. You can also modify the
+    text or use your own. Please note that these two dimensions can be
+    entangled, as a negative sentiment may also detoxify the text.
+    '''
+    if st.session_state.get("output", "") != "" and \
+            st.button("Analyze the styled text", type="primary"):
+        col1, col2 = st.columns(2)
+        for name, col, dim, color in zip(
+            ["Sentiment", "Detoxification"],
+            [col1, col2],
+            [2, 0],
+            ["#ff7f0e", "#1f77b4"],
+        ):
+            col.subheader(name)
+            # classification
+            col.markdown("##### Dimension-Wise Classification Distribution")
+            _, dist_list, _ = model.steer_analysis(
+                analyzed_text,
+                dim, -steer_range, steer_range,
+                bins=2*int(steer_range)+1,
+            )
+            dist_list = np.array(dist_list)
+            col.bar_chart(
+                pd.DataFrame(
+                    {
+                        "Value": dist_list[:, 0],
+                        "Probability": dist_list[:, 1],
+                    }
+                ), x="Value", y="Probability",
+                color=color,
+            )
+
+            # key tokens
+            pos_steer, neg_steer = np.zeros((2, 4))
+            pos_steer[dim] = 1
+            neg_steer[dim] = -1
+            _, token_evidence = model.evidence_words(
+                analyzed_text,
+                [pos_steer, neg_steer],
+            )
+            tokens = tokenizer(analyzed_text).input_ids
+            tokens = [f"{i:3d}: {tokenizer.decode([t])}"
+                      for i, t in enumerate(tokens)]
+            col.markdown("##### Token's Evidence Score in the Dimension")
+            col.bar_chart(
+                pd.DataFrame(
+                    {
+                        "Token": tokens[1:],
+                        "Evidence": token_evidence,
+                    }
+                ), x="Token", y="Evidence",
+                horizontal=True, color=color,
+            )
+
+    st.divider()
+    st.divider()
+    st.subheader("The Word Embeddings Space Analysis")
+    '''
+    LM-Steer provides a lens on how word embeddings correlate with LM word
+    embeddings: what word dimensions contribute to or contrast to a specific
+    style. This analysis can be used to understand the word embedding space
+    and how it steers the model's generation.
+    Note that due to the bidirectional nature of the embedding spaces, in each
+    dimension, sometimes only one side of the word embeddings is most relevant
+    to the style (can be either left or right).
+    '''
+    dimension = st.selectbox(
+        "Select a dimension to analyze",
+        ["Sentiment", "Detoxification"],
+    )
+    dim = 2 if dimension == "Sentiment" else 0
+    word_embedding_space_analysis(model, tokenizer, dim)
+
+
+if __name__ == "__main__":
+    main()

lm_steer/arguments.py

@@ -0,0 +1,59 @@
+from pprint import pprint
+import argparse
+from .utils import set_seed
+
+
+def parse_args():
+    # Model related
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str,
+                        default='EleutherAI/gpt-neo-2.7B')
+    parser.add_argument("--adaptor_class", type=str, default="multiply")
+    parser.add_argument("--adapted_component", type=str, default="final_layer")
+    parser.add_argument("--epsilon", type=float, default=1e-3)
+    parser.add_argument("--init_var", type=float, default=1e-2)
+    parser.add_argument("--rank", type=int, default=1000)
+    parser.add_argument("--num_steers", type=int, default=10)
+    parser.add_argument("--temperature", type=int, default=1)
+    parser.add_argument("--cuda", action="store_true")
+    parser.add_argument("--low_resource_mode", action="store_true")
+
+    # Data related
+    parser.add_argument("--data_dir", type=str, default=None)
+    parser.add_argument("--dataset_name", type=str, default=None)
+    parser.add_argument("--eval_file", type=str, default=None)
+    parser.add_argument("--output_file", type=str, default=None)
+    parser.add_argument("--data_size", type=int, default=None)
+    parser.add_argument("--split", type=str, default=None)
+
+    # Training related
+    parser.add_argument("--regularization", type=float, default=0)
+    parser.add_argument("--optimizer", type=str, default="Adam")
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--gamma_mean", type=float, default=0.99)
+    parser.add_argument("--n_steps", type=int, default=10000)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--ckpt_name", type=str, default=None)
+    parser.add_argument("--max_length", type=int, default=256)
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--log_step", type=int, default=500)
+    parser.add_argument("--subset", type=int, default=None)
+    parser.add_argument("--dummy_steer", type=int, default=None)
+    parser.add_argument("--training_steer", type=int, default=0)
+
+    # Evaluation related
+    parser.add_argument("--eval_size", type=int, default=None)
+    parser.add_argument("--steer_values", default=None, nargs="*", type=float)
+    parser.add_argument("--verbose", action="store_true")
+    parser.add_argument("--top_p", type=float, default=1)
+
+    # transfer related
+    parser.add_argument("--transfer_from", type=str, default=None)
+
+    args = parser.parse_args()
+
+    set_seed(args.seed)
+
+    print("arguments:")
+    pprint(args.__dict__)
+    return args

lm_steer/models/get_model.py

@@ -0,0 +1,43 @@
+
+def get_model(model_name, adapted_component, adaptor_class, num_steers, rank,
+              epsilon, init_var, low_resource_mode):
+    if model_name.startswith("EleutherAI/gpt-neo") or \
+            model_name.startswith("gpt2"):
+        from lm_steer.models.model_gpt_neo import Switching_GPTNeoModel
+        model = Switching_GPTNeoModel(
+            model_name, adapted_component, adaptor_class, num_steers, rank,
+            epsilon, init_var, low_resource_mode)
+        return model, model.tokenizer
+    elif model_name.startswith("lora-gpt2"):
+        from lm_steer.models.model_lora_gpt_neo import LORA_GPTNeoModel
+        model = LORA_GPTNeoModel(model_name, rank, epsilon)
+        return model, model.tokenizer
+    elif model_name.startswith("embedding_tuning"):
+        from lm_steer.models.model_embedding_tuning_gpt_neo import \
+            EmbeddingTuning_GPTNeoModel
+        model = EmbeddingTuning_GPTNeoModel(model_name)
+        return model, model.tokenizer
+    elif model_name.startswith("prefix-gpt2"):
+        from lm_steer.models.model_prefix_gpt_neo import PREFIX_GPTNeoModel
+        model = PREFIX_GPTNeoModel(model_name)
+        return model, model.tokenizer
+    elif model_name.startswith("EleutherAI/pythia"):
+        from lm_steer.models.model_gpt_neox import Switching_GPTNeoXModel
+        model = Switching_GPTNeoXModel(
+            model_name, adapted_component, adaptor_class, num_steers, rank,
+            epsilon, init_var, low_resource_mode)
+        return model, model.tokenizer
+    elif model_name.startswith("EleutherAI/gpt-j"):
+        from lm_steer.models.model_gpt_j import Switching_GPTJModel
+        model = Switching_GPTJModel(
+            model_name, adapted_component, adaptor_class, num_steers, rank,
+            epsilon, init_var, low_resource_mode)
+        return model, model.tokenizer
+    elif model_name.startswith("microsoft/DialoGPT"):
+        from lm_steer.models.model_dialogpt import Switching_DialoGPTModel
+        model = Switching_DialoGPTModel(
+            model_name, adapted_component, adaptor_class, num_steers, rank,
+            epsilon, init_var, low_resource_mode)
+        return model, model.tokenizer
+    else:
+        raise NotImplementedError()

lm_steer/models/model_base.py

@@ -0,0 +1,173 @@
+
+import torch.nn as nn
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+from lm_steer.utils import set_seed
+from .model_utils import find_max_subspans
+
+
+punctuations = [
+    '!', '"', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.',
+    # '/', '#',
+    ':', ';', '<', '=', '>', '?', '@',
+    '[', '\\', ']', '^', '_', '`',
+    '{', '|', '}', '~',
+    '¨', '©', 'ª', '«', '¬', '®', '¯', '°', '±', '²', '³', '´', 'µ', '¶', '·',
+    '¸', '¹', 'º', '»', '¼', '½', '¾',
+    '\n', ' ',
+]
+
+
+class LMSteerBase(nn.Module):
+    def evidence_words(self, prompt, comparing_steer_values,
+                       truncation_length=1024, max_segments=4, max_length=10):
+        if isinstance(comparing_steer_values, list):
+            comparing_steer_values = \
+                torch.Tensor(comparing_steer_values).to(self.device)
+        if (comparing_steer_values[0] - comparing_steer_values[1]
+                ).abs().sum() <= 0.2:
+            return [(prompt, None)]
+        tokenized = self.tokenizer(
+            prompt, return_tensors="pt",
+            max_length=truncation_length, truncation=True)
+        input_ids = torch.LongTensor(tokenized["input_ids"]).to(self.device)
+        input_ids = input_ids.expand(2, -1)
+        attention_mask = torch.LongTensor(tokenized["attention_mask"]).to(
+            self.device)
+        attention_mask = attention_mask.expand(2, -1)
+        self.steer.set_value(comparing_steer_values)
+        with torch.no_grad():
+            output = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=input_ids)
+        length = input_ids.shape[1]
+        loss_token = F.cross_entropy(
+            output.logits[:, :-1].reshape((2)*(length-1), -1),
+            input_ids[:, 1:].reshape(-1),
+            reduction="none"
+        )
+        loss_token = loss_token.reshape(2, length - 1)
+
+        token_evidence = (- loss_token[0] + loss_token[1])
+        tokens = input_ids[0]
+        evidence_segments = find_max_subspans(
+            token_evidence.cpu().numpy().tolist(), max_segments, max_length)[0]
+        evidence_segments = [
+            (_seg[0]+1, _seg[1]+1) for _seg in evidence_segments]
+        start = 0
+        output = []
+        if len(evidence_segments) > 0:
+            for _segment in evidence_segments:
+                if _segment[0] > start:
+                    output.append((
+                        self.tokenizer.decode(tokens[start: _segment[0]]),
+                        None
+                    ))
+                output.append((
+                    self.tokenizer.decode(tokens[_segment[0]: _segment[1]]),
+                    "evidence"
+                ))
+                start = _segment[1]
+            length = tokens.shape[-1]
+            if _segment[1] < length:
+                output.append((
+                    self.tokenizer.decode(tokens[_segment[1]: length]),
+                    None
+                ))
+        else:
+            output = [(prompt, None)]
+
+        return output, token_evidence.tolist()
+
+    def steer_analysis(self, prompt, steer_dim, min_value=-3, max_value=3,
+                       bins=7):
+        tokenized = self.tokenizer(prompt)
+        input_ids = torch.LongTensor(tokenized["input_ids"]).to(self.device)
+        input_ids = input_ids.expand(bins + 1, -1)
+        attention_mask = torch.LongTensor(tokenized["attention_mask"]).to(
+            self.device)
+        attention_mask = attention_mask.expand(bins + 1, -1)
+        steer_values = torch.zeros(bins+1, self.num_steers).to(self.device)
+        for bin_i in range(bins):
+            steer_values[bin_i, steer_dim] = (
+                min_value + (max_value - min_value) / (bins - 1) * bin_i
+            )
+        self.steer.set_value(steer_values)
+        with torch.no_grad():
+            output = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=input_ids)
+        length = input_ids.shape[1]
+        loss_token = F.cross_entropy(
+            output.logits[:, :-1].reshape((bins+1)*(length-1), -1),
+            input_ids[:, 1:].reshape(-1),
+            reduction="none"
+        )
+        loss_token = loss_token.reshape(bins + 1, length - 1)
+        loss = loss_token.mean(-1)[:-1]
+        dist = ((- loss + loss.mean()) * 100).softmax(0)
+        dist_list = list(zip(
+            [
+                min_value + (max_value - min_value) / (bins - 1) * bin_i
+                for bin_i in range(bins)
+            ],
+            dist.tolist(),
+        ))
+        best_guess = loss.argmin(0)
+        best_guess_value = min_value + \
+            (max_value - min_value) / (bins - 1) * best_guess.item()
+
+        token_evidence = (- loss_token[best_guess] + loss_token[-1]) * 10
+        token_evidence = [0] + token_evidence.tolist()
+        # tokens = self.tokenizer.convert_ids_to_tokens(input_ids[0])
+
+        word_evidence_list = []
+        start = 0
+        n_tokens = len(input_ids[0])
+        for token_i in range(1, n_tokens+1):
+            span = self.tokenizer.decode(input_ids[0][start: token_i])
+            for _punc in punctuations:
+                if token_i == n_tokens or _punc in span:
+                    new_span = self.tokenizer.decode(
+                        input_ids[0][start: token_i-1]).strip()
+                    if len(new_span) <= 1:
+                        break
+                    word_evidence_list.append((
+                        new_span,
+                        np.array(token_evidence[start: token_i-1]).mean()
+                    ))
+                    start = token_i - 1
+                    break
+
+        # token_evidence_list = list(zip(tokens, token_evidence))
+        return best_guess_value, dist_list, word_evidence_list
+
+    def generate(self, prompt, steer_values, min_length=20, max_length=100,
+                 seed=None, num_beams=1, num_beam_groups=1, do_sample=True,
+                 temperature=1, top_p=1):
+        '''
+        prompt: a string
+        steer_values
+        min_length: minimum generation length
+        max_length: maximum generation length
+        seed: seed for generation. None if not specified.
+        '''
+        if seed is not None:
+            set_seed(seed)
+        steer_values = torch.Tensor(steer_values).to(
+            self.device)
+        self.steer.set_value(steer_values[None])
+        with torch.no_grad():
+            text = self.generator(
+                prompt, num_beams=num_beams, num_beam_groups=num_beam_groups,
+                do_sample=do_sample, temperature=temperature, top_p=top_p,
+                min_length=min_length, max_length=max_length,
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+            text = text[0]["generated_text"]
+
+        return text

lm_steer/models/model_embedding_tuning_gpt_neo.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+from transformers import pipeline
+
+from .model_utils import Hack_no_grad
+from lm_steer.utils import set_seed
+
+
+class EmbeddingTuning_GPTNeoModel(nn.Module):
+    def __init__(self, model_name):
+        super().__init__()
+        self.generator = pipeline(
+            'text-generation',
+            model=model_name.replace("embedding_tuning-", ""))
+        self.tokenizer = self.generator.tokenizer
+        self.model = self.generator.model
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+        self.model.transformer = Hack_no_grad(self.model.transformer)
+
+    def forward(self, input_ids, attention_mask, steer_values):
+        output = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=input_ids)
+        return output
+
+    def parameters(self):
+        return [self.model.lm_head.weight]
+
+    def state_dict(self):
+        return self.model.lm_head.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.model.lm_head.load_state_dict(state_dict)
+
+    def to_device(self, device):
+        self.generator.device = device
+        self.model.to(device)
+        self.device = device
+
+    def regularization_term(self):
+        return torch.tensor(0)
+
+    def generate(self, prompt, steer_values, min_length=20, max_length=100,
+                 seed=None, num_beams=1, num_beam_groups=1, do_sample=True,
+                 temperature=1, top_p=1):
+        if seed is not None:
+            set_seed(seed)
+        with torch.no_grad():
+            text = self.generator(
+                prompt, num_beams=num_beams, num_beam_groups=num_beam_groups,
+                do_sample=do_sample, temperature=temperature, top_p=top_p,
+                min_length=min_length, max_length=max_length,
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+            text = text[0]["generated_text"]
+        return text

lm_steer/models/model_gpt_j.py

@@ -0,0 +1,333 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GPTJForCausalLM, AutoTokenizer
+
+from .model_utils import Hack_no_grad, find_max_subspans
+from .steers import Projected_Adaptor
+from lm_steer.utils import set_seed
+
+
+punctuations = [
+    '!', '"', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.',
+    # '/', '#',
+    ':', ';', '<', '=', '>', '?', '@',
+    '[', '\\', ']', '^', '_', '`',
+    '{', '|', '}', '~',
+    '¨', '©', 'ª', '«', '¬', '®', '¯', '°', '±', '²', '³', '´', 'µ', '¶', '·',
+    '¸', '¹', 'º', '»', '¼', '½', '¾',
+    '\n', ' ',
+]
+
+
+class Switching_GPTJModel(nn.Module):
+    def __init__(self, model_name, adapted_component, adaptor_class,
+                 num_steers, rank, epsilon, init_var, low_resource_mode):
+        super().__init__()
+        self.adapted_component = adapted_component
+        self.adaptor_class = adaptor_class
+        # self.generator = pipeline('text-generation', model=model_name)
+        # self.tokenizer = self.generator.tokenizer
+        # self.model = self.generator.model
+        if low_resource_mode:
+            print("using low_resource_mode and fp16")
+            self.model = GPTJForCausalLM.from_pretrained(
+                "EleutherAI/gpt-j-6B", revision="float16",
+                torch_dtype=torch.float16, low_cpu_mem_usage=True
+            )
+        else:
+            self.model = GPTJForCausalLM.from_pretrained(
+                "EleutherAI/gpt-j-6B",
+            )
+        self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.init_var = init_var
+        self.num_steers = num_steers
+        self.device = torch.device("cpu")
+        self.low_resource_mode = low_resource_mode
+        embed_dim = self.model.lm_head.weight.shape[1]
+        vocab_size = self.model.lm_head.weight.shape[0]
+
+        for _param in self.model.parameters():
+            _param.requires_grad_(False)
+
+        if adapted_component == "final_layer":
+            self.model.transformer = Hack_no_grad(self.model.transformer)
+            self.steer = Projected_Adaptor(
+                self.model.lm_head, adaptor_class, num_steers, embed_dim,
+                vocab_size, rank, epsilon, init_var, "output")
+            self.model.set_output_embeddings(self.steer)
+        elif adapted_component == "input_embedding":
+            self.steer = Projected_Adaptor(
+                self.model.transformer.wte, adaptor_class, num_steers,
+                embed_dim, vocab_size, rank, epsilon, init_var, "input")
+            self.model.transformer.set_input_embeddings(self.steer)
+        else:
+            raise NotImplementedError()
+
+    def forward(self, input_ids, attention_mask, steer_values):
+        self.steer.set_value(steer_values)
+        output = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=input_ids)
+        return output
+
+    def parameters(self):
+        return self.steer.parameters()
+
+    def state_dict(self):
+        return self.steer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.steer.load_state_dict(state_dict)
+
+    def to_device(self, device):
+        # self.generator.device = device
+        self.model.to(device)
+        self.device = device
+
+    def regularization_term(self):
+        return self.steer.regularization_term()
+
+    def generate(self, prompt, steer_values, min_length=20, max_length=100,
+                 seed=None, num_beams=1, num_beam_groups=1, do_sample=True,
+                 temperature=1, top_p=1):
+        '''
+        prompt: a string
+        steer_values
+        min_length: minimum generation length
+        max_length: maximum generation length
+        seed: seed for generation. None if not specified.
+        '''
+        if seed is not None:
+            set_seed(seed)
+        steer_values = torch.Tensor(steer_values).to(
+            self.device)
+        if self.low_resource_mode:
+            fp16 = torch.float16
+            steer_values = steer_values.to(fp16)
+            self.steer.projector1.data = self.steer.projector1.to(fp16)
+            self.steer.projector2.data = self.steer.projector2.to(fp16)
+        self.steer.set_value(steer_values[None])
+        with torch.no_grad():
+            input_ids = self.tokenizer(
+                prompt, return_tensors="pt").input_ids.to(self.device)
+            gen_tokens = self.model.generate(
+                input_ids,
+                num_beams=num_beams, num_beam_groups=num_beam_groups,
+                do_sample=do_sample, temperature=temperature, top_p=top_p,
+                min_new_tokens=min_length, max_new_tokens=max_length,
+                pad_token_id=self.tokenizer.pad_token_id)
+            text = self.tokenizer.batch_decode(gen_tokens)[0]
+
+        # recovering
+        if self.low_resource_mode:
+            fp32 = torch.float32
+            self.steer.projector1.data = self.steer.projector1.to(fp32)
+            self.steer.projector2.data = self.steer.projector2.to(fp32)
+        return text
+
+    def generate_multiple(
+            self, prompts, steer_values, min_length=20, max_length=100,
+            seed=None):
+        '''
+        prompt: a string
+        steer_values
+        min_length: minimum generation length
+        max_length: maximum generation length
+        seed: seed for generation. None if not specified.
+        '''
+        if seed is not None:
+            set_seed(seed)
+        steer_values = torch.Tensor(steer_values).to(
+            self.device)
+        if self.low_resource_mode:
+            fp16 = torch.float16
+            steer_values = steer_values.to(fp16)
+            self.steer.projector1.data = self.steer.projector1.to(fp16)
+            self.steer.projector2.data = self.steer.projector2.to(fp16)
+        self.steer.set_value(steer_values)
+        with torch.no_grad():
+            input_ids = self.tokenizer(
+                prompts, return_tensors="pt").input_ids.to(self.device)
+            gen_tokens = self.model.generate(
+                input_ids,
+                do_sample=True,
+                min_new_tokens=min_length, max_new_tokens=max_length,
+                pad_token_id=self.tokenizer.pad_token_id)
+            text = self.tokenizer.batch_decode(gen_tokens)
+
+        # recovering
+        if self.low_resource_mode:
+            fp32 = torch.float32
+            self.steer.projector1.data = self.steer.projector1.to(fp32)
+            self.steer.projector2.data = self.steer.projector2.to(fp32)
+        return text
+
+    # def evidence_words(self, prompt, original_steer_values, max_segments=4,
+    #                    max_length=10):
+    #     if isinstance(original_steer_values, list):
+    #         original_steer_values = torch.Tensor(original_steer_values)
+    #     if original_steer_values.abs().sum() <= 0.2:
+    #         return [(prompt, None)]
+    #     tokenized = self.tokenizer(prompt)
+    #     input_ids = torch.LongTensor(tokenized["input_ids"]).to(self.device)
+    #     input_ids = input_ids.expand(2, -1)
+    #     attention_mask = torch.LongTensor(tokenized["attention_mask"]).to(
+    #         self.device)
+    #     attention_mask = attention_mask.expand(2, -1)
+    #     steer_values = torch.zeros(2, self.num_steers).to(self.device)
+    #     steer_values[0] = original_steer_values
+    #     steer_values[1] = (-original_steer_values > 0) * 2 - 1
+    #     if self.low_resource_mode:
+    #         fp16 = torch.float16
+    #         steer_values = steer_values.to(fp16)
+    #         self.steer.projector1.data = self.steer.projector1.to(fp16)
+    #         self.steer.projector2.data = self.steer.projector2.to(fp16)
+    #     self.steer.set_value(steer_values)
+    #     with torch.no_grad():
+    #         output = self.model(
+    #             input_ids=input_ids,
+    #             attention_mask=attention_mask,
+    #             labels=input_ids)
+    #     length = input_ids.shape[1]
+    #     loss_token = F.cross_entropy(
+    #         output.logits[:, :-1].reshape((2)*(length-1), -1),
+    #         input_ids[:, 1:].reshape(-1),
+    #         reduction="none"
+    #     )
+    #     loss_token = loss_token.reshape(2, length - 1)
+
+    def evidence_words(self, prompt, original_steer_values,
+                       truncation_length=1024, max_segments=4, max_length=10):
+        if isinstance(original_steer_values, list):
+            original_steer_values = torch.Tensor(original_steer_values)
+        if original_steer_values.abs().sum() <= 0.2:
+            return [(prompt, None)]
+        tokenized = self.tokenizer(
+            prompt, return_tensors="pt", max_length=truncation_length, truncation=True)
+        input_ids = torch.LongTensor(tokenized["input_ids"]).to(self.device)
+        input_ids = input_ids.expand(2, -1)
+        attention_mask = torch.LongTensor(tokenized["attention_mask"]).to(
+            self.device)
+        attention_mask = attention_mask.expand(2, -1)
+        steer_values = torch.zeros(2, self.num_steers).to(self.device)
+        steer_values[0] = original_steer_values
+        steer_values[1] = (-original_steer_values > 0) * 2 - 1
+        if self.low_resource_mode:
+            fp16 = torch.float16
+            steer_values = steer_values.to(fp16)
+            self.steer.projector1.data = self.steer.projector1.to(fp16)
+            self.steer.projector2.data = self.steer.projector2.to(fp16)
+        self.steer.set_value(steer_values)
+        with torch.no_grad():
+            output = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=input_ids)
+        length = input_ids.shape[1]
+        loss_token = F.cross_entropy(
+            output.logits[:, :-1].reshape((2)*(length-1), -1),
+            input_ids[:, 1:].reshape(-1),
+            reduction="none"
+        )
+        loss_token = loss_token.reshape(2, length - 1)
+
+        token_evidence = (- loss_token[0] + loss_token[1])
+        tokens = input_ids[0]
+        evidence_segments = find_max_subspans(
+            token_evidence.cpu().numpy().tolist(), max_segments, max_length)[0]
+        evidence_segments = [
+            (_seg[0]+1, _seg[1]+1) for _seg in evidence_segments]
+        start = 0
+        output = []
+        color = (
+            "gray" if original_steer_values.shape[0] > 1
+            else "red" if original_steer_values[0] > 0
+            else "blue"
+        )
+        if len(evidence_segments) > 0:
+            for _segment in evidence_segments:
+                if _segment[0] > start:
+                    output.append((
+                        self.tokenizer.decode(tokens[start: _segment[0]]),
+                        None
+                    ))
+                output.append((
+                    self.tokenizer.decode(tokens[_segment[0]: _segment[1]]),
+                    color
+                ))
+                start = _segment[1]
+            length = tokens.shape[-1]
+            if _segment[1] < length:
+                output.append((
+                    self.tokenizer.decode(tokens[_segment[1]: length]),
+                    None
+                ))
+        else:
+            output = [(prompt, None)]
+
+        if self.low_resource_mode:
+            fp32 = torch.float32
+            self.steer.projector1.data = self.steer.projector1.to(fp32)
+            self.steer.projector2.data = self.steer.projector2.to(fp32)
+        return output
+
+    def steer_analysis(self, prompt, steer_dim, min_value=-3, max_value=3,
+                        bins=7, truncation_length=1024):
+        tokenized = self.tokenizer(
+            prompt, return_tensors="pt",
+            max_length=truncation_length,
+            truncation=True)
+        input_ids = torch.LongTensor(tokenized["input_ids"]).to(self.device)
+        input_ids = input_ids.expand(bins + 1, -1)
+        attention_mask = torch.LongTensor(tokenized["attention_mask"]).to(
+            self.device)
+        attention_mask = attention_mask.expand(bins + 1, -1)
+        steer_values = torch.zeros(bins+1, self.num_steers).to(self.device)
+        for bin_i in range(bins):
+            steer_values[bin_i, steer_dim] = (
+                min_value + (max_value - min_value) / (bins - 1) * bin_i
+            )
+        if self.low_resource_mode:
+            fp16 = torch.float16
+            steer_values = steer_values.to(fp16)
+            self.steer.projector1.data = self.steer.projector1.to(fp16)
+            self.steer.projector2.data = self.steer.projector2.to(fp16)
+        self.steer.set_value(steer_values)
+        with torch.no_grad():
+            output = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=input_ids)
+        length = input_ids.shape[1]
+        loss_token = F.cross_entropy(
+            output.logits[:, :-1].reshape((bins+1)*(length-1), -1),
+            input_ids[:, 1:].reshape(-1),
+            reduction="none"
+        )
+        loss_token = loss_token.reshape(bins + 1, length - 1)
+        loss = loss_token.mean(-1)[:-1]
+        dist = ((- loss + loss.mean()) * 100).softmax(0)
+        dist_list = list(zip(
+            [
+                min_value + (max_value - min_value) / (bins - 1) * bin_i
+                for bin_i in range(bins)
+            ],
+            dist.tolist(),
+        ))
+        best_guess = loss.argmin(0)
+        best_guess_value = min_value + \
+            (max_value - min_value) / (bins - 1) * best_guess.item()
+
+        token_evidence = self.evidence_words(
+            prompt, steer_values[best_guess],
+        )
+
+        if self.low_resource_mode:
+            fp32 = torch.float32
+            self.steer.projector1.data = self.steer.projector1.to(fp32)
+        return best_guess_value, dist_list, token_evidence

lm_steer/models/model_gpt_neo.py

@@ -0,0 +1,66 @@
+import torch
+from transformers import pipeline
+
+from .model_utils import Hack_no_grad
+from .steers import Projected_Adaptor
+from .model_base import LMSteerBase
+
+
+class Switching_GPTNeoModel(LMSteerBase):
+    def __init__(self, model_name, adapted_component, adaptor_class,
+                 num_steers, rank, epsilon, init_var,
+                 low_resource_mode):
+        super().__init__()
+        self.adapted_component = adapted_component
+        self.generator = pipeline('text-generation', model=model_name)
+        self.tokenizer = self.generator.tokenizer
+        self.model = self.generator.model
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.init_var = init_var
+        self.num_steers = num_steers
+        self.device = torch.device("cpu")
+        embed_dim = self.model.lm_head.weight.shape[1]
+        vocab_size = self.model.lm_head.weight.shape[0]
+
+        for _param in self.model.parameters():
+            _param.requires_grad_(False)
+
+        if adapted_component == "final_layer":
+            self.model.transformer = Hack_no_grad(self.model.transformer)
+            self.steer = Projected_Adaptor(
+                self.model.lm_head, adaptor_class, num_steers, embed_dim,
+                vocab_size, rank, epsilon, init_var, "output")
+            self.model.set_output_embeddings(self.steer)
+        elif adapted_component == "input_embedding":
+            self.steer = Projected_Adaptor(
+                self.model.transformer.wte, adaptor_class, num_steers,
+                embed_dim, vocab_size, rank, epsilon, init_var, "input")
+            self.model.transformer.set_input_embeddings(self.steer)
+        else:
+            raise NotImplementedError()
+
+    def forward(self, input_ids, attention_mask, steer_values):
+        self.steer.set_value(steer_values)
+        output = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=input_ids)
+        return output
+
+    def parameters(self):
+        return self.steer.parameters()
+
+    def state_dict(self):
+        return self.steer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.steer.load_state_dict(state_dict)
+
+    def to_device(self, device):
+        self.generator.device = device
+        self.model.to(device)
+        self.device = device
+
+    def regularization_term(self):
+        return self.steer.regularization_term()

lm_steer/models/model_gpt_neox.py

@@ -0,0 +1,105 @@
+import torch
+from transformers import GPTNeoXForCausalLM, AutoTokenizer
+
+from .model_utils import Hack_no_grad
+from .steers import Projected_Adaptor
+from .model_base import LMSteerBase
+from lm_steer.utils import set_seed
+
+
+class Switching_GPTNeoXModel(LMSteerBase):
+    def __init__(self, model_name, adapted_component, adaptor_class,
+                 num_steers, rank, epsilon, init_var,
+                 low_resource_mode):
+        super().__init__()
+        self.adapted_component = adapted_component
+        if low_resource_mode:
+            self.model = GPTNeoXForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16, low_cpu_mem_usage=True
+            )
+        else:
+            self.model = GPTNeoXForCausalLM.from_pretrained(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.init_var = init_var
+        self.num_steers = num_steers
+        self.device = torch.device("cpu")
+        embed_dim = self.model.embed_out.weight.shape[1]
+        vocab_size = self.model.embed_out.weight.shape[0]
+        self.low_resource_mode = low_resource_mode
+
+        for _param in self.model.parameters():
+            _param.requires_grad_(False)
+
+        if adapted_component == "final_layer":
+            self.model.gpt_neox = Hack_no_grad(self.model.gpt_neox)
+            self.steer = Projected_Adaptor(
+                self.model.embed_out, adaptor_class, num_steers, embed_dim,
+                vocab_size, rank, epsilon, init_var, "output")
+            self.model.set_output_embeddings(self.steer)
+        else:
+            raise NotImplementedError()
+
+    def forward(self, input_ids, attention_mask, steer_values):
+        self.steer.set_value(steer_values)
+        output = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=input_ids)
+        return output
+
+    def parameters(self):
+        return self.steer.parameters()
+
+    def state_dict(self):
+        return self.steer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.steer.load_state_dict(state_dict)
+
+    def to_device(self, device):
+        self.model.to(device)
+        self.device = device
+
+    def regularization_term(self):
+        return self.steer.regularization_term()
+
+    def generate(self, prompt, steer_values, min_length=20, max_length=100,
+                 seed=None, num_beams=1, num_beam_groups=1, do_sample=True,
+                 temperature=1, top_p=1):
+        '''
+        prompt: a string
+        steer_values
+        min_length: minimum generation length
+        max_length: maximum generation length
+        seed: seed for generation. None if not specified.
+        '''
+        if seed is not None:
+            set_seed(seed)
+        steer_values = torch.Tensor(steer_values).to(
+            self.device)
+        if self.low_resource_mode:
+            fp16 = torch.float16
+            steer_values = steer_values.to(fp16)
+            self.steer.projector1.data = self.steer.projector1.to(fp16)
+            self.steer.projector2.data = self.steer.projector2.to(fp16)
+        self.steer.set_value(steer_values[None])
+        with torch.no_grad():
+            input_ids = self.tokenizer(
+                prompt, return_tensors="pt").input_ids.to(self.device)
+            gen_tokens = self.model.generate(
+                input_ids,
+                num_beams=num_beams, num_beam_groups=num_beam_groups,
+                do_sample=do_sample, temperature=temperature, top_p=top_p,
+                min_length=min_length, max_length=max_length,
+                pad_token_id=self.tokenizer.pad_token_id)
+            text = self.tokenizer.batch_decode(gen_tokens)[0]
+
+        # recovering
+        if self.low_resource_mode:
+            fp32 = torch.float32
+            self.steer.projector1.data = self.steer.projector1.to(fp32)
+            self.steer.projector2.data = self.steer.projector2.to(fp32)
+        return text

lm_steer/models/model_lora_gpt_neo.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+from transformers import pipeline
+from peft import LoraConfig, get_peft_model
+
+from lm_steer.utils import set_seed
+
+
+class LORA_GPTNeoModel(nn.Module):
+    def __init__(self, model_name, rank, epsilon):
+        super().__init__()
+        self.generator = pipeline('text-generation',
+                                  model=model_name.replace("lora-", ""))
+        self.tokenizer = self.generator.tokenizer
+        model = self.generator.model
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+        config = LoraConfig(
+            r=rank,
+            lora_alpha=epsilon,
+            target_modules=["c_attn", "c_proj", "c_fc"],
+            lora_dropout=0.1,
+            bias="lora_only",
+            modules_to_save=[],
+        )
+        self.model = get_peft_model(model, config)
+        self.generator.model = self.model
+        self.model.print_trainable_parameters()
+
+    def forward(self, input_ids, attention_mask, steer_values):
+        output = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=input_ids)
+        return output
+
+    def to_device(self, device):
+        self.generator.device = device
+        self.model.to(device)
+        self.device = device
+
+    def regularization_term(self):
+        return torch.tensor(0)
+
+    def generate(self, prompt, steer_values, min_length=20, max_length=100,
+                 seed=None, num_beams=1, num_beam_groups=1, do_sample=True,
+                 temperature=1, top_p=1):
+        if seed is not None:
+            set_seed(seed)
+        with torch.no_grad():
+            text = self.generator(
+                prompt, num_beams=num_beams, num_beam_groups=num_beam_groups,
+                do_sample=do_sample, temperature=temperature, top_p=top_p,
+                min_length=min_length, max_length=max_length,
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+            text = text[0]["generated_text"]
+        return text

lm_steer/models/model_utils.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class Hack_no_grad(nn.Module):
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+
+    def forward(self, *inputs, **kwargs):
+        with torch.no_grad():
+            return self.module(*inputs, **kwargs)
+
+
+def find_max_subspans(sequence, n_spans, max_length):
+    length = len(sequence)
+    inner_scores = np.zeros((length, n_spans + 1, max_length + 1, 2))
+    trace = np.zeros((length, n_spans + 1, max_length + 1, 2, 3), dtype=int)
+    # trace[:, n_spans, max_length, 0] = (n_spans, max_length, 0)
+    inner_scores[-1, :, :, 1] = -1e5
+    for _i in range(length):
+        for _j in range(n_spans+1):
+            for _k in range(max_length+1):
+                trace[_i, _j, _k, 0] = (_j, max_length, 0)
+
+    for _i in range(length):
+        for _j in range(n_spans):
+            for _k in range(max_length+1):
+                inner_scores[_i, _j, _k, 0], trace[_i, _j, _k, 0] = (
+                    inner_scores[_i-1, _j, max_length, 0],
+                    (_j, max_length, 0)
+                )
+                max_taken = inner_scores[_i-1, _j, :, 1].max()
+                if max_taken > inner_scores[_i, _j, _k, 0]:
+                    inner_scores[_i, _j, _k, 0] = max_taken
+                    trace[_i, _j, _k, 0] = (
+                        _j, inner_scores[_i-1, _j, :, 1].argmax(), 1)
+
+                if _k < max_length:
+                    inner_scores[_i, _j, _k, 1], trace[_i, _j, _k, 1] = (
+                        (
+                            inner_scores[_i-1, _j, _k+1, 1] + sequence[_i],
+                            (_j, _k+1, 1)
+                        )
+                        if (inner_scores[_i-1, _j, _k+1, 1] >
+                            inner_scores[_i-1, _j+1, max_length, 0])
+                        else (
+                            inner_scores[_i-1, _j+1, max_length, 0] +
+                            sequence[_i],
+                            (_j+1, max_length, 0)
+                        )
+                    )
+
+    max_score = 0
+    argmax = (0, 0, 0)
+    for _j in reversed(range(n_spans + 1)):
+        for _k in reversed(range(max_length)):
+            if inner_scores[-1, _j, _k, 0] > max_score:
+                max_score = inner_scores[-1, _j, _k, 0]
+                argmax = (_j, _k, 0)
+            if inner_scores[-1, _j, _k, 1] > max_score:
+                max_score = inner_scores[-1, _j, _k, 1]
+                argmax = (_j, _k, 1)
+
+    trace_back = argmax
+    tags = []
+    for _i in reversed(range(length)):
+        tags.append(trace_back[2])
+        trace_back = trace[_i, trace_back[0], trace_back[1], trace_back[2]]
+
+    tags.reverse()
+    segments = []
+    start = None
+    for _i in range(length + 1):
+        if _i < length and tags[_i] == 1 and start is None:
+            start = _i
+        elif (_i == length or tags[_i] == 0) and start is not None:
+            segments.append((start, _i))
+            start = None
+    return segments, max_score, tags  # , inner_scores, trace

lm_steer/models/steers.py

@@ -0,0 +1,96 @@
+import torch
+import torch.nn as nn
+
+
+class Projected_Adaptor(nn.Module):
+    def __init__(self, lm_head, adaptor_class, num_steers, embed_dim,
+                 vocab_size, rank, epsilon, init_var, position="output"):
+        super().__init__()
+        assert rank > 0
+        if adaptor_class == "multiply":
+            self.projector1 = nn.Parameter(torch.randn(
+                num_steers, embed_dim, rank
+            ) * init_var)
+            self.projector2 = nn.Parameter(torch.randn(
+                num_steers, embed_dim, rank
+            ) * init_var)
+        elif adaptor_class == "add":
+            self.add_vec = nn.Parameter(torch.randn(
+                num_steers, embed_dim
+            ))
+        elif adaptor_class == "offset":
+            self.offset_vec = nn.Parameter(torch.randn(
+                num_steers, vocab_size
+            ))
+        else:
+            raise NotImplementedError()
+
+        self.adaptor_class = adaptor_class
+        self.rank = rank
+        self.lm_head = lm_head
+        self.epsilon = epsilon
+        self.position = position
+        self.num_steers = num_steers
+        self.init_var = init_var
+        self.steer_values = torch.zeros(num_steers)
+
+    def set_value(self, steer_values):
+        self.steer_values = steer_values
+
+    def forward(self, state):
+        if self.steer_values.abs().sum() == 0:
+            return state.matmul(
+                self.lm_head.weight.detach().transpose(0, 1))
+        if self.adaptor_class == "multiply":
+            delta = state[:, None].matmul(self.projector1[None]) *\
+                self.steer_values[:, :, None, None]
+            delta = delta.matmul(
+                self.projector2.transpose(1, 2)[None]).sum(1)
+            projected_state = state + self.epsilon * delta
+            logits = projected_state.matmul(
+                self.lm_head.weight.detach().transpose(0, 1))
+        elif self.adaptor_class == "add":
+            add_values = self.steer_values.matmul(self.add_vec)
+            projected_state = state + self.epsilon * add_values[:, None]
+            logits = projected_state.matmul(
+                self.lm_head.weight.detach().transpose(0, 1))
+        elif self.adaptor_class == "offset":
+            offset_values = self.steer_values.matmul(self.offset_vec)
+            logits = state.matmul(
+                self.lm_head.weight.detach().transpose(0, 1))
+            logits = logits + self.epsilon * offset_values[:, None]
+        return logits
+
+    def regularization_term(self):
+        if self.adaptor_class == "multiply":
+            return self.projector1.pow(2).sum() + self.projector2.pow(2).sum()
+        elif self.adaptor_class == "add":
+            return self.add_vec.pow(2).sum()
+        elif self.adaptor_class == "offset":
+            return self.offset_vec.pow(2).sum()
+
+    def parameters(self):
+        if self.adaptor_class == "multiply":
+            return [self.projector1, self.projector2]
+        elif self.adaptor_class == "add":
+            return [self.add_vec]
+        elif self.adaptor_class == "offset":
+            return [self.offset_vec]
+
+    def state_dict(self):
+        if self.adaptor_class == "multiply":
+            return {"projector1": self.projector1,
+                    "projector2": self.projector2}
+        elif self.adaptor_class == "add":
+            return {"add_vec": self.add_vec}
+        elif self.adaptor_class == "offset":
+            return {"offset_vec": self.offset_vec}
+
+    def load_state_dict(self, state_dict):
+        if self.adaptor_class == "multiply":
+            self.projector1.data = state_dict["projector1"]
+            self.projector2.data = state_dict["projector2"]
+        elif self.adaptor_class == "add":
+            self.add_vec.data = state_dict["add_vec"]
+        elif self.adaptor_class == "offset":
+            self.offset_vec.data = state_dict["offset_vec"]

lm_steer/utils.py

@@ -0,0 +1,45 @@
+import random
+import torch
+import numpy as np
+
+
+def set_seed(seed):
+    if seed is None:
+        return
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+class RunningMean:
+    def __init__(self, gamma):
+        self.gamma = gamma
+        self.count = 0
+        self._value = None
+
+    def update(self, value):
+        value = value.detach().cpu()
+        if value.ndim == 0:
+            self._update(value)
+        else:
+            for _v in value:
+                self._update(_v)
+
+    def _update(self, value):
+        self.count += 1
+        if self._value is None:
+            self._value = value
+        else:
+            w1 = self.gamma * (1 - self.gamma ** (self.count - 1))
+            w2 = (1 - self.gamma)
+            wt = w1 + w2
+            w1 = w1 / wt
+            w2 = w2 / wt
+            self._value = w1 * self._value + w2 * value
+
+    @property
+    def value(self):
+        if self._value is None:
+            return 0
+        return self._value * 1

requirements.txt

@@ -0,0 +1,5 @@
+torch
+transformers
+datasets
+numpy
+pandas