app.py

# 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 (the larger the more positive)",
        -steer_range, steer_range, 3.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)
    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"):
        with st.spinner("Generating..."):
            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"],
        ):
            with st.spinner(f"Analyzing {name}..."):
                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.write("The polarity of the token's evidence score "
                          "which aligns with sliding bar directions."
                          )
                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
    with st.spinner("Analyzing..."):
        word_embedding_space_analysis(model, tokenizer, dim)


if __name__ == "__main__":
    main()