Change to local inference

activations/candidate_vectors.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed63186d01ddaf6df8835818144185b5fb05d1c9a4683fce9517a921472353b3
+size 804046

activations/deepseek-1.5b-candidate_vectors.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a4f3701085a9090e78fc402aaaef5adbf23f0b49c932f82eb4fc107d191aac0
+size 345294

activations/deepseek-1.5b-offsets.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d84cb880bee5feb83248b476d8d0f3f87dca74bc8ae53807f3ab2a9bdb959920
+size 345244

activations/offsets.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6cb10bd9014f9cd2470d37f56f491abd5f72bd162543a7569f40cd385f127c3
+size 803996

app.py

@@ -1,35 +1,25 @@
-import os, json
+import threading
 import logging
 from pathlib import Path
-import asyncio
-import aiohttp
+from typing import Dict
+import spaces
 import pandas as pd
 import gradio as gr
 from gradio_toggle import Toggle
+from transformers import TextIteratorStreamer
+from model import load_model
 from scheduler import load_scheduler
 from schemas import UserRequest, SteeringOutput, CONFIG
 
-
-MAX_RETRIES = 10
-MAX_RETRY_WAIT_TIME = 75
-MIN_RETRY_WAIT_TIME = 5
-ENDPOINT_ALIVE = False
-
-HF_TOKEN = os.getenv('HF_TOKEN')
-API_URL = "https://a6k5m81qw14hkvhz.us-east-1.aws.endpoints.huggingface.cloud"
-headers = {
-    "Accept" : "application/json",
-    "Authorization": f"Bearer {HF_TOKEN}",
-    "Content-Type": "application/json" 
-}
-
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(name)s %(levelname)s:%(message)s')
 logger = logging.getLogger(__name__)
 
 model_name = "DeepSeek-R1-Distill-Qwen-7B"
 examples = pd.read_csv("assets/examples.csv")
+
 instances = {}
 scheduler = load_scheduler()
+model = load_model()
 
 
 HEAD = """
@@ -198,8 +188,6 @@ def initialize_instance(request: gr.Request):
 
 
 def cleanup_instance(request: gr.Request):
-    global ENDPOINT_ALIVE
-
     session_id = request.session_hash
 
     if session_id in instances:
@@ -209,51 +197,48 @@ def cleanup_instance(request: gr.Request):
                 
         del instances[session_id]
 
-    if len(instances) == 0:
-        ENDPOINT_ALIVE = False
-
     logger.info("Number of connections: %d", len(instances))
 
 
-async def initialize_endpoint():
-    alive = False
-    session = aiohttp.ClientSession()
-    async with session.get(f"{API_URL}/health", headers=headers) as resp:
-        resp_text = await resp.text()
-        if resp.status == 200:
-            alive = True
-        else:
-            logger.error("API Error Code: %d, Message: %s", resp.status, resp_text)
-    
-    await session.close()
-    return alive
+@spaces.GPU(duration=90)
+def generate(prompt: str, steering: bool, coeff: float, generation_config: Dict[str, float], layer: int, k: float):
+    formatted_prompt = model.apply_chat_template(prompt)
+    inputs = model.tokenize(formatted_prompt)
 
+    streamer = TextIteratorStreamer(model.tokenizer, timeout=10, skip_prompt=True, skip_special_tokens=True)
 
-async def get_endpoint_state():
-    global ENDPOINT_ALIVE
-    n = 0
-    sleep_time = MAX_RETRY_WAIT_TIME
-
-    while n < MAX_RETRIES:
-        n += 1
+    if steering:
+        thread = threading.Thread(
+            target=model.steer_generation,
+            args=(inputs, streamer, k, layer, coeff, generation_config)
+        )
+    else:
+        thread = threading.Thread(
+            target=model.run_generation,
+            args=(inputs, streamer, generation_config)
+        )
 
-        if not ENDPOINT_ALIVE:
-            logger.info("Initializing inference endpoint")
-            yield "Initializing"
-            ENDPOINT_ALIVE = await initialize_endpoint()
+    thread.start()
+    
+    generated_text = "<think>"
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text
+
+
+def generate_output(
+    session_id: str, prompt: str, steering: bool, coeff: float, 
+    max_new_tokens: int, top_p: float, temperature: float, layer: int, vec_scaling: float
+):
+    req = UserRequest(
+        session_id=session_id, prompt=prompt, steering=steering, coeff=coeff,
+        max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, vec_scale=vec_scaling, layer=layer
+    )
 
-        if ENDPOINT_ALIVE:
-            logger.info("Inference endpoint is ready")
-            gr.Info("Inference endpoint is ready")
-            yield "Ready"
-            break
-        
-        gr.Warning("Initializing inference endpoint\n(This may take 2~3 minutes)", duration=sleep_time)
-        await asyncio.sleep(sleep_time)
-        sleep_time = max(sleep_time * 0.8, MIN_RETRY_WAIT_TIME)
+    logger.info("User request: %s", req)
+    instances[session_id].append(req)
 
-    if n == MAX_RETRIES:
-        yield "Server Error"
+    yield from generate(prompt, steering, coeff, req.generation_config(), layer, req.k)
 
 
 async def post_process(session_id, output):
@@ -266,62 +251,11 @@ async def post_process(session_id, output):
             answer = None
         else:
             answer = p[-1]
-    else:
-        answer = None
-        reasoning = output
-
-    steering_output = SteeringOutput(**req.model_dump(), reasoning=reasoning, answer=answer)
-    instances[session_id].append(steering_output)
-
-
-class Generator:
-    def __init__(self):
-        self.stop_events = {}
-
-    async def stop(self, session_id):
-        self.stop_events[session_id] = True
-        logger.info("Stopping generation")
-
-    async def generate(
-        self, session_id: str, prompt: str, steering: bool, coeff: float, 
-        max_new_tokens: int, top_p: float, temperature: float, layer: int, vec_scaling: float
-    ):
-        req = UserRequest(
-            session_id=session_id, prompt=prompt, steering=steering, coeff=coeff,
-            max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, vec_scale=vec_scaling, layer=layer
-        )
-
-        instances[session_id].append(req)
-
-        data = req.get_api_format()
-        logger.info("User Request: %s", data)
-
-        generated_text = ""
-        self.stop_events[session_id] = False
-
-        try:
-            async with aiohttp.ClientSession() as session:
-                async with session.post(f"{API_URL}/generate", headers=headers, json=data) as resp:
-                    if resp.status == 200:
-                        generated_text += "<think>"
-
-                        async for chunk, _ in resp.content.iter_chunks():
-                            if self.stop_events[session_id]:
-                                break
 
-                            generated_text += chunk.decode()
-                            yield generated_text
-                    else:
-                        logger.error("API Error Ccode: %d, Error Message: %s", resp.status, resp.text())
-                        raise gr.Error("API Server Error")
-                    
-        except:
-            logger.info("Client session error")
+        steering_output = SteeringOutput(**req.model_dump(), reasoning=reasoning, answer=answer)
+        instances[session_id].append(steering_output)
 
-        if generated_text != "":
-            await post_process(session_id, generated_text)
-
-        del self.stop_events[session_id]
+    return gr.update(interactive=True), gr.update(interactive=True)
 
 
 async def output_feedback(session_id, feedback):
@@ -339,31 +273,13 @@ async def output_feedback(session_id, feedback):
         logger.debug("Feedback submission error")
 
 
-async def show_feedback_buttons(upvote_btn, downvote_btn):
-    return gr.update(interactive=True), gr.update(interactive=True)
-
-
 gr.set_static_paths(paths=[Path.cwd().absolute() / "assets"])
 theme = gr.themes.Base(primary_hue="emerald", text_size=gr.themes.sizes.text_lg).set()
-generator = Generator()
 
 with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS, js=JS) as demo:
     session_id = gr.State()
-    endpoint_state = gr.State(get_endpoint_state)
-
     gr.HTML(HTML)
-
-    @gr.render(inputs=endpoint_state, triggers=[endpoint_state.change])
-    def render_state(endpoint_state):
-        if endpoint_state == "Ready":
-            color = "green"
-        elif endpoint_state == "Server Error":
-            color = "red"
-        else:
-            color = "orange"
-
-        if endpoint_state != None:
-            gr.Markdown(f'🤖 {model_name} | Inference Endpoint State: <span style="color:{color}; font-weight: bold;">{endpoint_state}</span>', elem_id="model-state")
+    gr.Markdown(f'🤖 {model_name}')
 
     with gr.Row(elem_id="main-components"):
         with gr.Column(scale=1):
@@ -382,7 +298,6 @@ with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS,
 
             with gr.Row():
                 clear_btn = gr.ClearButton()
-                stop_btn = gr.Button("Stop")
                 generate_btn = gr.Button("Generate", variant="primary")
 
             with gr.Accordion("⚙️ Advanced Settings", open=False):
@@ -408,25 +323,19 @@ with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS,
     gr.Examples(examples=examples[examples["type"] == "harmful"].prompt.tolist(), inputs=input_text, label="Harmful")
 
 
-    @gr.on(triggers=[clear_btn.click, stop_btn.click], outputs=[upvote_btn, downvote_btn])
-    def clear_feedback_buttons():
+    @gr.on(triggers=[clear_btn.click], outputs=[upvote_btn, downvote_btn])
+    def clear():
         return gr.update(interactive=False), gr.update(interactive=False)
-    
-    @gr.on(triggers=[generate_btn.click], outputs=[upvote_btn, downvote_btn])
-    def show_feedback_buttons():
-        return gr.update(interactive=True), gr.update(interactive=True)
-
-
-    submission = generate_btn.click(
-        generator.generate, inputs=[session_id, input_text, steer_toggle, coeff, max_new_tokens, top_p, temperature, layer, vec_scaling], outputs=output
-    )
 
     clear_btn.add([input_text, output])
-    stop_btn.click(generator.stop, inputs=session_id, queue=False)
+    generate_btn.click(
+        generate_output, inputs=[session_id, input_text, steer_toggle, coeff, max_new_tokens, top_p, temperature, layer, vec_scaling], outputs=output
+    ).success(
+        post_process, inputs=[session_id, output], outputs=[upvote_btn, downvote_btn]
+    )
 
     upvote_btn.click(output_feedback, inputs=[session_id, upvote_btn])
     downvote_btn.click(output_feedback, inputs=[session_id, downvote_btn])
-
     layer.change(fn=lambda x: 1, inputs=vec_scaling, outputs=vec_scaling)
 
     demo.load(initialize_instance, outputs=session_id)

model.py

@@ -0,0 +1,110 @@
+import os, warnings
+from operator import attrgetter
+from typing import List, Dict
+
+import torch
+import torch.nn.functional as F
+from torchtyping import TensorType
+from transformers import TextIteratorStreamer
+from transformers import AutoTokenizer, BatchEncoding
+import nnsight
+from nnsight import LanguageModel
+from nnsight.intervention import Envoy
+
+warnings.filterwarnings("ignore")
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# nnsight with multi-threading: https://github.com/ndif-team/nnsight/issues/280
+nnsight.CONFIG.APP.GLOBAL_TRACING = False
+
+config = {
+    "model_name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
+    "steering_vec": "activations/candidate_vectors.pt",
+    "offset": "activations/offsets.pt",
+}
+
+def detect_module_attrs(model: LanguageModel) -> str:
+    if "model" in model._modules and "layers" in model.model._modules:
+        return "model.layers"
+    elif "transformers" in model._modules and "h" in model.transformers._modules:
+        return "transformers.h"
+    else:
+        raise Exception("Failed to detect module attributes.")
+
+
+class ModelBase:
+    def __init__(
+        self, model_name: str,
+        steering_vecs: TensorType, offsets: TensorType,
+        tokenizer: AutoTokenizer = None, block_module_attr=None
+    ):
+        if tokenizer is None:
+            self.tokenizer = self._load_tokenizer(model_name)
+        else:
+            self.tokenizer = tokenizer
+        self.model = self._load_model(model_name, self.tokenizer)
+
+        self.device = self.model.device
+        self.hidden_size = self.model.config.hidden_size
+        if block_module_attr is None:
+            self.block_modules = self.get_module(detect_module_attrs(self.model))
+        else:
+            self.block_modules = self.get_module(block_module_attr)
+
+        self.steering_vecs = F.normalize(steering_vecs, dim=-1)
+        self.steering_vecs, self.offsets = self.set_dtype(self.steering_vecs, offsets)
+    
+    def _load_model(self, model_name: str, tokenizer: AutoTokenizer) -> LanguageModel:
+        return LanguageModel(model_name, tokenizer=tokenizer, dispatch=True, trust_remote_code=True, device_map="auto", torch_dtype=torch.bfloat16)
+    
+    def _load_tokenizer(self, model_name) -> AutoTokenizer:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer.padding_side = "left"
+        if not tokenizer.pad_token:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+            tokenizer.pad_token = tokenizer.eos_token
+
+        tokenizer.chat_template = tokenizer.chat_template.replace("<｜Assistant｜><think>\\n", "<｜Assistant｜><think>")
+        return tokenizer
+
+    def tokenize(self, prompt: str) -> BatchEncoding:
+        return self.tokenizer(prompt, padding=True, truncation=False, return_tensors="pt")
+    
+    def get_module(self, attr: str) -> Envoy:
+        return attrgetter(attr)(self.model)
+
+    def set_dtype(self, *vars):
+        if len(vars) == 1:
+            return vars[0].to(self.model.dtype)
+        else:
+            return (var.to(self.model.dtype) for var in vars)
+    
+    def apply_chat_template(self, instruction: str) -> List[str]:
+        messages = [{"role": "user", "content": instruction}]
+        return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    def run_generation(self, inputs, streamer: TextIteratorStreamer, generation_config: Dict):
+        inputs = inputs.to(self.device)
+        _ = self.model._model.generate(**inputs, do_sample=True, streamer=streamer, **generation_config)
+
+    def steer_generation(
+        self, inputs, streamer: TextIteratorStreamer, k: float, 
+        layer: int, coeff: float, generation_config: Dict
+    ):
+        layer_block = self.block_modules[layer]
+        unit_vec = self.steering_vecs[layer]
+        offset = self.offsets[layer]
+
+        with self.model.generate(inputs, do_sample=True, streamer=streamer, **generation_config):
+            with self.block_modules.all():
+                acts = layer_block.output[0].clone()
+                proj = (acts - offset) @ unit_vec.unsqueeze(-1) * unit_vec
+                layer_block.output[0][:] = acts - proj + coeff * k * unit_vec
+
+
+def load_model() -> ModelBase:
+    steering_vecs = torch.load(config['steering_vec'], weights_only=True)
+    offsets = torch.load(config['offset'], weights_only=True)
+    model = ModelBase(config['model_name'], steering_vecs=steering_vecs, offsets=offsets)
+    return model
+

requirements.txt

@@ -2,3 +2,11 @@ aiohttp==3.11.16
 pandas==2.2.2
 pyarrow==19.0.1
 gradio_toggle==2.0.2
+transformers==4.50.0
+accelerate==1.6.0
+nnsight==0.4.3
+triton==3.1.0
+torchtyping==0.1.5
+tiktoken==0.8.0
+transformers_stream_generator==0.0.5
+zstandard==0.23.0

scheduler.py

@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 
 def load_scheduler():
     return ParquetScheduler(
-        repo_id="hannahcyberey/Censorship-Steering-Logs", every=10,
+        repo_id="hannahcyberey/Censorship-Steering-Logs", every=60,
         private=True,
         squash_history=False,
         schema={

schemas.py

@@ -32,18 +32,11 @@ class UserRequest(BaseModel):
         else:
             self.k = self.vec_scale * vector_scaling[self.layer]["k_neg"]
 
-    def get_api_format(self):
+    def generation_config(self):
         return {
-            "prompt": self.prompt,
-            "steering": self.steering,
-            "coeff": self.coeff,
-            "k": self.k,
-            "layer": self.layer,
-            "generation_config": {
-                "max_new_tokens": self.max_new_tokens,
-                "top_p": self.top_p,
-                "temperature": self.temperature
-            }
+            "max_new_tokens": self.max_new_tokens,
+            "top_p": self.top_p,
+            "temperature": self.temperature
         }