app.py

import os, json, random
import torch
import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from huggingface_hub import login, hf_hub_download
import pyreft
import pyvene as pv
from threading import Thread
from typing import Iterator

HF_TOKEN = os.environ.get("HF_TOKEN")
login(token=HF_TOKEN)

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 128  # smaller default to save memory
MAX_INPUT_TOKEN_LENGTH = 4096

def load_jsonl(jsonl_path):
    jsonl_data = []
    with open(jsonl_path, 'r') as f:
        for line in f:
            data = json.loads(line)
            jsonl_data.append(data)
    return jsonl_data

class Steer(pv.SourcelessIntervention):
    """Steer model via activation addition"""
    def __init__(self, **kwargs):
        super().__init__(**kwargs, keep_last_dim=True)
        self.proj = torch.nn.Linear(
            self.embed_dim, kwargs["latent_dim"], bias=False
        )
    def forward(self, base, source=None, subspaces=None):
        steering_vec = torch.tensor(subspaces["mag"]) * \
            self.proj.weight[subspaces["idx"]].unsqueeze(dim=0)
        return base + steering_vec

# Check GPU
if not torch.cuda.is_available():
    print("Warning: Running on CPU, may be slow.")

# Load model & dictionary
model_id = "google/gemma-2-2b-it"
pv_model = None
tokenizer = None
concept_list = []
concept_id_map = {}
if torch.cuda.is_available():
    model = AutoModelForCausalLM.from_pretrained(
        model_id, device_map="cuda", torch_dtype=torch.bfloat16
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id)

    # Download dictionary
    weight_path = hf_hub_download(repo_id="pyvene/gemma-reft-2b-it-res", filename="l20/weight.pt")
    meta_path = hf_hub_download(repo_id="pyvene/gemma-reft-2b-it-res", filename="l20/metadata.jsonl")
    params = torch.load(weight_path).cuda()
    md = load_jsonl(meta_path)

    concept_list = [item["concept"] for item in md]
    concept_id_map = {item["concept"]: item["concept_id"] for item in md}

    steer = Steer(embed_dim=params.shape[0], latent_dim=params.shape[1])
    steer.proj.weight.data = params.float()

    pv_model = pv.IntervenableModel({
        "component": f"model.layers[20].output",
        "intervention": steer}, model=model)

terminators = [tokenizer.eos_token_id] if tokenizer else []

@spaces.GPU
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    subspaces_list: list[dict],
    max_new_tokens: int=DEFAULT_MAX_NEW_TOKENS,
) -> Iterator[str]:
    # limit to last 3 turns
    start_idx = max(0, len(chat_history) - 3)
    recent_history = chat_history[start_idx:]

    # build list of messages
    messages = []
    # for user_msg, model_msg in recent_history:
    #     messages.append({"role": "user", "content": user_msg})
    #     messages.append({"role": "model", "content": model_msg})
    messages.append({"role": "user", "content": message})

    input_ids = torch.tensor([tokenizer.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True)]).cuda()

    # trim if needed
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        yield "[Truncated prior text]\n"

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = {
        "base": {"input_ids": input_ids},
        "unit_locations": None,
        "max_new_tokens": max_new_tokens,
        "intervene_on_prompt": True,
        "subspaces": [
            {
                "idx": int(subspaces_list[0]["idx"]),
                "mag": int(subspaces_list[0]["internal_mag"])
            }
        ] if subspaces_list else [],
        "streamer": streamer,
        "do_sample": True
    }

    t = Thread(target=pv_model.generate, kwargs=generate_kwargs)
    t.start()

    partial_text = []
    for token_str in streamer:
        partial_text.append(token_str)
        yield "".join(partial_text)

def filter_concepts(search_text: str):
    if not search_text.strip():
        return concept_list[:500]
    filtered = [c for c in concept_list if search_text.lower() in c.lower()]
    return filtered[:500]

def add_concept_to_list(selected_concept, user_slider_val, current_list):
    """
    Return exactly 2 values:
      1) The updated list of concepts (list of dicts).
      2) A Gradio update for the removal dropdown’s choices.
    """
    if not selected_concept:
        return current_list, gr.update(choices=_build_remove_choices(current_list))

    idx = concept_id_map[selected_concept]
    internal_mag = user_slider_val * 50
    new_entry = {
        "text": selected_concept,
        "idx": idx,
        "display_mag": user_slider_val,
        "internal_mag": internal_mag,
    }
    updated_list = current_list + [new_entry]
    return updated_list, gr.update(choices=_build_remove_choices(updated_list))

def remove_concept_from_list(selected_text, current_list):
    """
    Return exactly 2 values:
      1) The updated list of concepts (list of dicts).
      2) A Gradio update for the removal dropdown’s choices.
    """
    if not selected_text:
        return current_list, gr.update(choices=_build_remove_choices(current_list))
    updated_list = [x for x in current_list if x["text"] != selected_text]
    return updated_list, gr.update(choices=_build_remove_choices(updated_list))

def _build_remove_choices(subspaces):
    return [x["text"] for x in subspaces]

def update_dropdown_choices(search_text):
    filtered = filter_concepts(search_text)
    return gr.update(choices=filtered)

with gr.Blocks(css="style.css") as demo:

    # Pre-populate with a random concept if available
    default_subspaces = []
    if pv_model and concept_list:
        default_concept = "words related to time travel and its consequences"
        default_subspaces = [{
            "text": default_concept,
            "idx": concept_id_map[default_concept],
            "display_mag": 3,
            "internal_mag": 150.0,
        }]

    selected_subspaces = gr.State(default_subspaces)
    with gr.Row():
        # Left side: bigger chat area
        with gr.Column(scale=7):
            chat_interface = gr.ChatInterface(
                fn=generate,
                title="LM Steering with ReFT-r1 (16K concepts)",
                type="messages",
                additional_inputs=[selected_subspaces],
            )
        # Right side: concept management
        with gr.Column(scale=3):
            gr.Markdown("# Steering Concepts")
            search_box = gr.Textbox(
                label="Search concepts",
                placeholder="e.g. 'time travel'"
            )
            concept_dropdown = gr.Dropdown(
                label="Filtered Concepts",
                choices=[]
            )
            concept_magnitude = gr.Slider(
                label="Steering Factor",
                minimum=-5,
                maximum=5,
                step=1,
                value=3
            )
            add_button = gr.Button("Add Concept")

            # Row with the remove dropdown + button
            with gr.Row():
                remove_dropdown = gr.Dropdown(
                    label="Remove concept",
                    choices=_build_remove_choices(default_subspaces),
                    multiselect=False
                )
                remove_button = gr.Button("Remove", variant="secondary")

    # Wire up events
    # When the search box changes, update the concept dropdown choices:
    search_box.change(
        update_dropdown_choices,
        [search_box],
        [concept_dropdown]
    )

    # When "Add Concept" is clicked, add the concept + magnitude to the list,
    # and update the "Remove" dropdown choices.
    add_button.click(
        add_concept_to_list,
        [concept_dropdown, concept_magnitude, selected_subspaces],
        [selected_subspaces, remove_dropdown]
    )

    # When "Remove" is clicked, remove the selected concept from the list,
    # and update the "Remove" dropdown choices.
    remove_button.click(
        remove_concept_from_list,
        [remove_dropdown, selected_subspaces],
        [selected_subspaces, remove_dropdown]
    )

    demo.launch()