caching

.gitignore

@@ -0,0 +1,166 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# repo-specific
+**/.DS_Store
+_logs
+_logs/
+checkpoints/

app.py

@@ -1,6 +1,7 @@
 # https://huggingface.co/spaces/Glaciohound/LM-Steer
 
 import torch
+import nltk
 import streamlit as st
 import random
 import numpy as np
@@ -23,39 +24,53 @@ def st_get_model(model_name, low_resource_mode):
     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)
+@st.cache_data()
+def word_embedding_space_analysis(
+        model_name, dim):
+    model = st.session_state.model
+    tokenizer = st.session_state.tokenizer
+    projector1 = model.steer.projector1.data[dim]
+    projector2 = model.steer.projector2.data[dim]
     embeddings = model.steer.lm_head.weight
+    matrix = projector1.matmul(projector2.transpose(0, 1))
+    S, V, D = torch.linalg.svd(matrix)
 
     data = []
-    for _i in range(10):
-        left_tokens = embeddings.matmul(D[_i]).argsort()[-20:].flip(0)
-        right_tokens = embeddings.matmul(D[_i]).argsort()[:20]
+    top = 30
+    select_words = 20
+    n_dim = 10
+    for _i in range(n_dim):
+        left_tokens = embeddings.matmul(D[_i]).argsort()[-top:].flip(0)
+        right_tokens = embeddings.matmul(D[_i]).argsort()[:top]
 
         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
+                if (
+                    len(word) > 2 and not word[0].isalpha() and
+                    word[1:].isalpha() and word[1:].lower().islower()
+                ):
+                    word = word[1:]
+                    if word.lower() in nltk.corpus.words.words():
+                        output.append(word)
+            return output[:select_words]
 
         data.append([
             ", ".join(filter_words(side_tokens))
             for side_tokens in [left_tokens, right_tokens]
         ])
-    st.table(pd.DataFrame(
+    return 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)
+    nltk.download('words')
     title = "LM-Steer: Word Embeddings Are Steers for Language Models"
     st.set_page_config(
         layout="wide",
@@ -92,14 +107,14 @@ def main():
     '''
     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.
+    https://github.com/Glaciohound/LM-Steer to host larger models.
     Some generated texts may contain toxic or offensive content. Please be
     cautious when using the generated texts.
     Note that for these smaller models, the generation quality may not be as
     good as the larger models (GPT-4, Llama, etc.).
     '''
     col1, col2 = st.columns(2)
-    st.session_state.model_name = col1.selectbox(
+    model_name = col1.selectbox(
         "Select a model to steer",
         [
             "gpt2",
@@ -108,48 +123,57 @@ def main():
             "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/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
+    # 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
+    low_resource_mode = False
     model, tokenizer = st_get_model(
-        st.session_state.model_name, low_resource_mode)
+        model_name, low_resource_mode)
+    st.session_state.model = model
+    st.session_state.tokenizer = tokenizer
     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))
+    st.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
+    steer_range = 3.
+    steer_interval = 0.2
     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 (the larger the more positive)",
-        -steer_range, steer_range, 3.0, steer_interval)
+        -steer_range, steer_range, 0.0, steer_interval)
     detoxification = col2.slider(
         "Detoxification Strength (the larger the less toxic)",
         -steer_range, steer_range, 0.0,
         steer_interval)
-    max_length = col3.number_input("Max length", 50, 300, 50, 50)
+    max_length = col3.number_input("Max length", 20, 200, 20, 20)
     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"):
+        if sentiment == 0 and detoxification == 0:
+            '''
+            **The steer values are both 0, which means the steered model
+            is the same as the original model.**
+            '''
         with st.spinner("Generating..."):
             steer_values = [detoxification, 0, sentiment, 0]
             st.session_state.output = model.generate(
@@ -159,8 +183,9 @@ def main():
                 min_length=0,
                 max_length=max_length,
                 do_sample=True,
+                top_p=0.9,
             )
-    analyzed_text = \
+    st.session_state.analyzed_text = \
         st.text_area("Generated text:", st.session_state.output, height=200)
 
     # Analysing the sentence
@@ -174,7 +199,7 @@ def main():
     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 \
+    if st.session_state.get("analyzed_text", "") != "" and \
             st.button("Analyze the styled text", type="primary"):
         col1, col2 = st.columns(2)
         for name, col, dim, color in zip(
@@ -187,9 +212,9 @@ def main():
                 col.subheader(name)
                 # classification
                 col.markdown(
-                    "##### Dimension-Wise Classification Distribution")
+                    "##### Sentence Classification Distribution")
                 _, dist_list, _ = model.steer_analysis(
-                    analyzed_text,
+                    st.session_state.analyzed_text,
                     dim, -steer_range, steer_range,
                     bins=2*int(steer_range)+1,
                 )
@@ -209,10 +234,10 @@ def main():
                 pos_steer[dim] = 1
                 neg_steer[dim] = -1
                 _, token_evidence = model.evidence_words(
-                    analyzed_text,
+                    st.session_state.analyzed_text,
                     [pos_steer, neg_steer],
                 )
-                tokens = tokenizer(analyzed_text).input_ids
+                tokens = tokenizer(st.session_state.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")
@@ -241,13 +266,13 @@ def main():
     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
-    with st.spinner("Analyzing..."):
-        word_embedding_space_analysis(model, tokenizer, dim)
+    for dimension in ["Sentiment", "Detoxification"]:
+        f'##### {dimension} Dimension'
+        dim = 2 if dimension == "Sentiment" else 0
+        analysis_result = word_embedding_space_analysis(
+            model_name, dim)
+        with st.expander("Show the analysis results"):
+            st.table(analysis_result)
 
 
 if __name__ == "__main__":

lm_steer/models/model_base.py

@@ -26,8 +26,8 @@ class LMSteerBase(nn.Module):
         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:
+        if (comparing_steer_values[0] - comparing_steer_values[1]).abs().sum()\
+                <= 0.2:
             return [(prompt, None)]
         tokenized = self.tokenizer(
             prompt, return_tensors="pt",
@@ -162,12 +162,77 @@ class LMSteerBase(nn.Module):
             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,
+            inputs = self.tokenizer(
+                prompt, return_tensors="pt").to(self.device)
+            text = self.model.generate(
+                **inputs,
+                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"]
+            text = self.tokenizer.decode(text[0], skip_special_tokens=True)
 
         return text
+
+    def generate_low_resource(
+        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)
+        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
+        fp32 = torch.float32
+        self.steer.projector1.data = self.steer.projector1.to(fp32)
+        self.steer.projector2.data = self.steer.projector2.to(fp32)
+        return text
+
+    def state_dict(self):
+        return self.steer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.steer.load_state_dict(state_dict)
+
+    def parameters(self):
+        return self.steer.parameters()
+
+    def to_device(self, device):
+        self.model.to(device)
+        self.device = device
+
+    def regularization_term(self):
+        return self.steer.regularization_term()
+
+    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

lm_steer/models/model_gpt_j.py

@@ -1,27 +1,14 @@
 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 .model_base import LMSteerBase
 from lm_steer.utils import set_seed
 
 
-punctuations = [
-    '!', '"', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.',
-    # '/', '#',
-    ':', ';', '<', '=', '>', '?', '@',
-    '[', '\\', ']', '^', '_', '`',
-    '{', '|', '}', '~',
-    '¨', '©', 'ª', '«', '¬', '®', '¯', '°', '±', '²', '³', '´', 'µ', '¶', '·',
-    '¸', '¹', 'º', '»', '¼', '½', '¾',
-    '\n', ' ',
-]
-
-
-class Switching_GPTJModel(nn.Module):
+class Switching_GPTJModel(LMSteerBase):
     def __init__(self, model_name, adapted_component, adaptor_class,
                  num_steers, rank, epsilon, init_var, low_resource_mode):
         super().__init__()
@@ -67,31 +54,6 @@ class Switching_GPTJModel(nn.Module):
         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):
@@ -102,33 +64,9 @@ class Switching_GPTJModel(nn.Module):
         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
+        return super().generate_low_resource(
+            prompt, steer_values, min_length, max_length, seed,
+            num_beams, num_beam_groups, do_sample, temperature, top_p)
 
     def generate_multiple(
             self, prompts, steer_values, min_length=20, max_length=100,
@@ -167,13 +105,14 @@ class Switching_GPTJModel(nn.Module):
             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):
+    # 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)
+    #     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(
@@ -201,133 +140,98 @@ class Switching_GPTJModel(nn.Module):
     #     )
     #     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 = (- 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)]
 
-        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)
+    #         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
+    #     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

@@ -1,6 +1,7 @@
 import torch
 from transformers import pipeline
 
+
 from .model_utils import Hack_no_grad
 from .steers import Projected_Adaptor
 from .model_base import LMSteerBase
@@ -12,9 +13,9 @@ class Switching_GPTNeoModel(LMSteerBase):
                  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.pipeline = pipeline('text-generation', model=model_name)
+        self.model = self.pipeline.model
+        self.tokenizer = self.pipeline.tokenizer
         self.tokenizer.pad_token = self.tokenizer.eos_token
         self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
         self.init_var = init_var
@@ -39,28 +40,3 @@ class Switching_GPTNeoModel(LMSteerBase):
             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

@@ -4,7 +4,6 @@ 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):
@@ -42,30 +41,6 @@ class Switching_GPTNeoXModel(LMSteerBase):
         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):
@@ -76,30 +51,6 @@ class Switching_GPTNeoXModel(LMSteerBase):
         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
+        return super().generate_low_resource(
+            prompt, steer_values, min_length, max_length, seed,
+            num_beams, num_beam_groups, do_sample, temperature, top_p)