add inference endpoint

activations/deepseek-r1-7b-offset.pt

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

activations/deepseek-r1-7b-steering-vec.pt

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

app.py

@@ -1,85 +1,94 @@
+import os
 import logging, json
-import threading
 from pathlib import Path
-from datetime import datetime, timezone
-import spaces
+import asyncio
+import aiohttp
 import pandas as pd
-from transformers import TextIteratorStreamer
 import gradio as gr
 from gradio_toggle import Toggle
-from model import load_model
-from scheduler import ParquetScheduler
+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__)
 
-scheduler = ParquetScheduler(
-    repo_id="hannahcyberey/Censorship-Steering-Logs", every=10,
-    private=True,
-    schema={
-        "prompt": {"_type": "Value", "dtype": "string"},
-        "steering": {"_type": "Value", "dtype": "bool"},
-        "coeff": {"_type": "Value", "dtype": "float64"},
-        "top_p": {"_type": "Value", "dtype": "float64"},
-        "temperature": {"_type": "Value", "dtype": "float64"},
-        "reasoning": {"_type": "Value", "dtype": "string"},
-        "answer": {"_type": "Value", "dtype": "string"},
-        "timestamp": {"_type": "Value", "dtype": "string"},
-    }
-)
-
-default_model = "DeepSeek-R1-Distill-Qwen-7B"
-model = load_model()
-default_config = {"max_new_tokens": 2048, "temperature": 0.6, "top_p": 0.95}
+model_name = "DeepSeek-R1-Distill-Qwen-7B"
 examples = pd.read_csv("assets/examples.csv")
+instances = {}
+scheduler = load_scheduler()
+
 
 HEAD = """
 <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.7.2/css/all.min.css" integrity="sha512-Evv84Mr4kqVGRNSgIGL/F/aIDqQb7xQ2vcrdIwxfjThSH8CSR7PBEakCr51Ck+w+/U6swU2Im1vVX0SVk9ABhg==" crossorigin="anonymous" referrerpolicy="no-referrer" />
 """
 
 HTML = f"""
-<div align="center" style="padding-bottom: var(--spacing-xl);">
-    <h1><img src="/gradio_api/file=assets/rudder_3094973.png"> LLM Censorship Steering </h1>
-    <div id="cover">
-        <img style="height: 120px;" src="/gradio_api/file=assets/demo-cover.png">
-        <p>🤖: {default_model}</p>
+<div id="banner">
+    <h1 style="font-size: 32px; line-height: 1.5em; margin-bottom: 0em;">
+        <img src="/gradio_api/file=assets/rudder_3094973.png" style="display: inline; height: 1.5em;"> LLM Censorship Steering 
+    </h1>
+    <div id="cover" style="height: 130px;">
+        <img style="height: 100%; padding-top: 0.5em;" src="/gradio_api/file=assets/demo-cover.png">
     </div>
 </div>
 """
 
 CSS = """
-h1 {font-size: 32px; line-height: 1.5em; margin-bottom: 0em;}
-img {display: inline; height: 1.5em;}
-a {font-size: 18px;}
-div#cover {
-    display: flex; 
-    flex-direction: column; 
-    align-items: flex-start; 
-    width: fit-content;
-    padding-top: 1em;
-}
-label span {color: var(--body-text-color);}
-.slider_input_container span {color: var(--body-text-color);}
-.slider_input_container {
+div#banner {
     display: flex;
-    flex-wrap: wrap;
-    input {appearance: auto;}
+    flex-direction: column;
+    align-items: center;
+    justify-content: center;
+}
+div#component-8 .form {
+    padding-top: 7.5px;
+    background: var(--block-background-fill);
+}
+div#component-9 {
+    .toggle-label {color: var(--body-text-color);}
+    span p {
+        font-size: var(--block-info-text-size);
+        line-height: var(--line-sm);
+        color: var(--block-label-text-color);
+    }
 }
-label span p {color: var(--block-label-text-color);}
-.toggle-label {color: var(--body-text-color);}
-div#component-4 .form {align-items: center; background-color: var(--block-background-fill);}
-div#component-6 {padding-bottom: 0;}
-div#component-6 .wrap .head {
+div#component-10 {
+    .slider_input_container span {color: var(--body-text-color);}
+    .slider_input_container {
+        display: flex;
+        flex-wrap: wrap;
+        input {appearance: auto;}
+    }
+}
+div#component-10 .wrap .head {
     justify-content: unset;
     label {margin-right: var(--size-2);}
-    label span {margin-bottom: 0;}
+    label span {
+        color: var(--body-text-color);
+        margin-bottom: 0;
+    }
 }
 """
 
 
 slider_info = """\
 <div style='display: flex; justify-content: space-between; line-height: normal;'>\
-    <span style='font-size: var(--block-info-text-size);'>Less censorship</span><span style='font-size: var(--block-info-text-size);'>More censorship</span>\
+    <span style='font-size: var(--block-info-text-size); color: var(--block-label-text-color);'>Less censorship</span>\
+    <span style='font-size: var(--block-info-text-size); color: var(--block-label-text-color);'>More censorship</span>\
 </div>\
 """\
 
@@ -109,113 +118,214 @@ async() => {
 """ % (slider_info, slider_ticks)
 
 
-class Generator():
-    def __init__(self):
-        self.data = {}
-    
-    @spaces.GPU(duration=90)
-    def generate_output(self, prompt, steering, coeff, generation_config):
-        streamer = TextIteratorStreamer(model.tokenizer, timeout=10, skip_prompt=True, skip_special_tokens=True)
-
-        thread = threading.Thread(
-            target=model.generate,
-            args=(prompt, streamer, steering, coeff, generation_config)
-        )
-        thread.start()
-        
-        generated_text = "<think>"
-        for new_text in streamer:
-            generated_text += new_text
-            yield generated_text
+def initialize_instance(request: gr.Request):
+    instances[request.session_hash] = []
+    logger.info("Number of connections: %d", len(instances))
+    return request.session_hash
 
-        thread.join()
-    
-    def run(
-        self, prompt: str, steering: bool, coeff: float, 
-        max_new_tokens: int, top_p: float = 1.0, temperature: float = 1.0
-    ):
-        self.data = {
-            "prompt": prompt,
-            "steering": steering,
-            "coeff": coeff,
-            "top_p": top_p,
-            "temperature": temperature,
-        }
-        generation_config = {
-            "max_new_tokens": max_new_tokens, 
-            "temperature": temperature, 
-            "top_p": top_p
-        }
-        logger.info("steering=%s, coeff=%0.1f, generation_config=%s", str(steering), coeff, repr(generation_config))
-
-        yield from self.generate_output(prompt, steering, coeff, generation_config)
-
-    def save_output(self, output: str):
-        if "</think>" in output:
-            p = [p for p in output.partition("</think>") if p != ""]
-            reasoning = "".join(p[:-1])
-            if len(p) == 1:
-                answer = None
-            else:
-                answer = p[-1]
-        else:
-            answer = None
-            reasoning = output
 
-        self.data["reasoning"] = reasoning
-        self.data["answer"] = answer
-        self.data["timestamp"] = datetime.now(timezone.utc).isoformat()
-        scheduler.append(self.data)
+def cleanup_instance(request: gr.Request):
+    global ENDPOINT_ALIVE
+
+    session_id = request.session_hash
 
+    if session_id in instances:
         with open("outputs.jsonl", "a") as f:
-            json.dump(self.data, f)
-            f.write("\n")
+            for data in instances[session_id]:
+                scheduler.append(data.model_dump())
+                json.dump(data.model_dump(), f)
+                f.write("\n")
+
+        del instances[session_id]
+
+    if len(instances) == 0:
+        ENDPOINT_ALIVE = False
+
+    logger.info("Number of connections: %d", len(instances))
+
+
+async def initialize_endpoint():
+    async with aiohttp.ClientSession() as session:
+        async with session.get(f"{API_URL}/health", headers=headers) as resp:
+            if resp.status == 200:
+                return True
+            else:
+                resp_text = await resp.text()
+                logger.error("API Error Code: %d, Message: %s", resp.status, resp_text)
+                return False
+
 
+async def get_endpoint_state():
+    global ENDPOINT_ALIVE
+    n = 0
+    sleep_time = MAX_RETRY_WAIT_TIME
 
-def steering_switch(toggle_value):
-    if toggle_value is True:
-        return gr.update(label="Steering"), gr.update(interactive=True)
+    while n < MAX_RETRIES:
+        n += 1
+
+        if not ENDPOINT_ALIVE:
+            logger.info("Initializing inference endpoint")
+            yield "Initializing"
+            ENDPOINT_ALIVE = await initialize_endpoint()
+
+        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)
+
+    if n == MAX_RETRIES:
+        yield "Server Error"
+
+
+async def save_output(req: UserRequest, output: str):
+    if "</think>" in output:
+        p = [p for p in output.partition("</think>") if p != ""]
+        reasoning = "".join(p[:-1])
+        if len(p) == 1:
+            answer = None
+        else:
+            answer = p[-1]
     else:
-        return gr.update(label="No Steering"), gr.update(interactive=False)
+        answer = None
+        reasoning = output
+
+    steering_output = SteeringOutput(**req.model_dump(), reasoning=reasoning, answer=answer)
+    instances[req.session_id].append(steering_output)
+
+
+async def generate(
+    session_id: str, prompt: str, steering: bool, coeff: float, 
+    max_new_tokens: int, top_p: float, temperature: 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
+    )
+
+    data = req.get_api_format()
+    logger.info("User Request: %s", data)
+
+    generated_text = ""
+    session = aiohttp.ClientSession()
+
+    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():
+                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")
+    
+    await session.close()
+
+    if generated_text != "":
+        await save_output(req, generated_text)
+
+
+async def post_process(session_id):
+    return instances[session_id][-1].request_id, gr.update(interactive=True), gr.update(interactive=True)
+
+
+async def output_feedback(session_id, request_id, feedback):
+    logger.info("Feedback received for request %s: %s", str(request_id), feedback)
+    try:
+        data = instances[session_id].pop()
+        if data.request_id == request_id:
+            if "Upvote" in feedback:
+                setattr(data, "upvote", True)
+            elif "Downvote" in feedback:
+                setattr(data, "upvote", False)
+
+        instances[session_id].append(data)
+        gr.Info("Thank you for your feedback!")
+    except:
+        logger.debug("Feedback submission error")
 
 
 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()
+    request_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>')
+
     with gr.Row():
         with gr.Column(scale=1):
             with gr.Row():
-                steer_toggle = Toggle(label="Steering", value=True, interactive=True, scale=0.2)
-                coeff = gr.Slider(label="Steering Coefficient:", value=-1, minimum=-2, maximum=2, step=0.1, scale=0.8, show_reset_button=False)
+                steer_toggle = Toggle(label="Steering", info="Turn off to generate original outputs", value=True, interactive=True, scale=2)
+                coeff = gr.Slider(label="Steering Coefficient:", value=-1.0, minimum=-2, maximum=2, step=0.1, scale=8, show_reset_button=False)
+
+            @gr.on(inputs=[steer_toggle], outputs=[steer_toggle, coeff], triggers=[steer_toggle.change])
+            def update_toggle(toggle_value):
+                if toggle_value is True:
+                    return gr.update(label="Steering", info="Turn off to generate original outputs"), gr.update(interactive=True)
+                else:
+                    return gr.update(label="No Steering", info="Turn on to steer model outputs"), gr.update(interactive=False)
 
             with gr.Accordion("⚙️ Advanced Settings", open=False):
                 with gr.Row():
-                    temperature = gr.Slider(0, 1, step=0.1, value=default_config['temperature'], interactive=True, label="Temperature", scale=1)
-                    top_p = gr.Slider(0, 1, step=0.1, value=default_config['top_p'], interactive=True, label="Top p", scale=1)
-                    max_new_tokens = gr.Number(minimum=10, maximum=2048, value=default_config['max_new_tokens'], interactive=True, label="Max new tokens", scale=0.5)
+                    temperature = gr.Slider(0, 1, step=0.1, value=CONFIG["temperature"], interactive=True, label="Temperature", scale=2)
+                    top_p = gr.Slider(0, 1, step=0.1, value=CONFIG["top_p"], interactive=True, label="Top p", scale=2)
+                    max_new_tokens = gr.Number(CONFIG["max_new_tokens"], minimum=10, maximum=CONFIG["max_new_tokens"], interactive=True, label="Max new tokens", scale=1)
             
-            input_text = gr.Textbox(label="Input", placeholder="Enter your prompt here...", lines=5)
+            input_text = gr.Textbox(label="Input", placeholder="Enter your prompt here...", lines=6, interactive=True)
 
             with gr.Row():
                 clear_btn = gr.ClearButton()
                 generate_btn = gr.Button("Generate", variant="primary")
 
         with gr.Column(scale=1):
-            output = gr.Textbox(label="Output", lines=16, max_lines=16)
+            output = gr.Textbox(label="Output", lines=15, max_lines=15, interactive=False)
+
+            with gr.Row():
+                upvote_btn = gr.Button("👍 Upvote", interactive=False)
+                downvote_btn = gr.Button("👎 Downvote", interactive=False)
 
     gr.HTML("<p>‼️ For research purposes, we log user inputs and generated outputs. Please avoid submitting any confidential or personal information.</p>")
     gr.Markdown("#### Examples")
     gr.Examples(examples=examples[examples["type"] == "sensitive"].prompt.tolist(), inputs=input_text, label="Sensitive")
     gr.Examples(examples=examples[examples["type"] == "harmful"].prompt.tolist(), inputs=input_text, label="Harmful")
-    
 
-    steer_toggle.change(steering_switch, inputs=steer_toggle, outputs=[steer_toggle, coeff])
-    generate_btn.click(generator.run, inputs=[input_text, steer_toggle, coeff, max_new_tokens, top_p, temperature], outputs=output).then(generator.save_output, inputs=output)
+    @gr.on(triggers=[clear_btn.click], outputs=[request_id, upvote_btn, downvote_btn])
+    def clear():
+        return None, gr.update(interactive=False), gr.update(interactive=False)
+
     clear_btn.add([input_text, output])
+    generate_btn.click(
+        generate, inputs=[session_id, input_text, steer_toggle, coeff, max_new_tokens, top_p, temperature], outputs=output
+    ).success(
+        post_process, inputs=session_id, outputs=[request_id, upvote_btn, downvote_btn]
+    )
+
+    upvote_btn.click(output_feedback, inputs=[session_id, request_id, upvote_btn])
+    downvote_btn.click(output_feedback, inputs=[session_id, request_id, downvote_btn])
+
+    demo.load(initialize_instance, outputs=session_id)
+    demo.unload(cleanup_instance)
     
+
 if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=5)
     demo.launch(debug=True)

model.py

@@ -1,118 +0,0 @@
-import os, warnings
-from operator import attrgetter
-from typing import List, Dict, Callable, Tuple
-
-import torch
-import torch.nn.functional as F
-from torchtyping import TensorType
-from transformers import TextIteratorStreamer
-from transformers import AutoTokenizer, BatchEncoding
-from nnsight import LanguageModel
-from nnsight.intervention import Envoy
-
-warnings.filterwarnings("ignore")
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-config = {
-    "model_name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
-    "steering_vec": "activations/deepseek-r1-7b-steering-vec.pt",
-    "offset": "activations/deepseek-r1-7b-offset.pt",
-    "layer": 25,
-    "k": 200,
-}
-
-
-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.")
-
-
-def orthogonal_projection(a: TensorType[..., -1], unit_vec: TensorType[-1]) -> TensorType[..., -1]:
-    return a @ unit_vec.unsqueeze(-1) * unit_vec
-
-
-def get_intervention_func(steering_vec: TensorType, offset=0, k=0, coeff=0) -> Callable:
-    """Get function for model intervention."""
-    unit_vec = F.normalize(steering_vec, dim=-1)
-    rescaled_vec = unit_vec * k
-    return lambda acts: acts - orthogonal_projection(acts - offset, unit_vec) + coeff * rescaled_vec
-
-
-
-class ModelBase:
-    def __init__(
-        self, model_name: str,
-        steering_vec: TensorType, offset: TensorType,
-        k: float, steering_layer: int,
-        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_layer = steering_layer
-        self.k = k
-        self.steering_vec, self.offset = self.set_dtype(steering_vec, offset)
-    
-    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
-        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 generate(self, prompt: str, streamer: TextIteratorStreamer, steering: bool, coeff: float, generation_config: Dict):
-        formatted_prompt = self.apply_chat_template(prompt)
-        inputs = self.tokenize(formatted_prompt)
-
-        if steering:
-            intervene_func = get_intervention_func(self.steering_vec, offset=self.offset, k=self.k, coeff=coeff)
-
-            with self.model.generate(inputs, do_sample=True, streamer=streamer, **generation_config):
-                self.block_modules.all()
-                acts = self.block_modules[self.steering_layer].output[0]
-                new_acts = intervene_func(acts)
-                self.block_modules[self.steering_layer].output[0][:] = new_acts
-        else:
-            inputs = inputs.to(self.device)
-            _ = self.model._model.generate(**inputs, do_sample=True, streamer=streamer, **generation_config)
-
-
-def load_model() -> ModelBase:
-    steering_vec = torch.load(config['steering_vec'], weights_only=True)
-    offset = torch.load(config['offset'], weights_only=True)
-    model = ModelBase(config['model_name'], steering_vec=steering_vec, offset=offset, k=config['k'], steering_layer=config['layer'])
-    model.tokenizer.chat_template = model.tokenizer.chat_template.replace("<｜Assistant｜><think>\\n", "<｜Assistant｜><think>")
-    return model

requirements.txt

@@ -1,11 +1,4 @@
-transformers==4.47.1
-accelerate==0.33.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
+aiohttp==3.11.16
 pandas==2.2.2
 pyarrow==19.0.1
 gradio_toggle==2.0.2

scheduler.py

@@ -2,7 +2,6 @@ import json
 import logging
 import tempfile
 import uuid
-from pathlib import Path
 from typing import Optional, Union, Dict, List, Any
 
 import pyarrow as pa
@@ -13,49 +12,38 @@ from huggingface_hub.hf_api import HfApi
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(name)s %(levelname)s:%(message)s')
 logger = logging.getLogger(__name__)
 
-
-def _infer_schema(key: str, value: Any) -> Dict[str, str]:
-    """
-    Infer schema for the `datasets` library.
-
-    See https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Value.
-    """
-    if "image" in key:
-        return {"_type": "Image"}
-    if "audio" in key:
-        return {"_type": "Audio"}
-    if isinstance(value, int):
-        return {"_type": "Value", "dtype": "int64"}
-    if isinstance(value, float):
-        return {"_type": "Value", "dtype": "float64"}
-    if isinstance(value, bool):
-        return {"_type": "Value", "dtype": "bool"}
-    if isinstance(value, bytes):
-        return {"_type": "Value", "dtype": "binary"}
-    # Otherwise in last resort => convert it to a string
-    return {"_type": "Value", "dtype": "string"}
+def load_scheduler():
+    return ParquetScheduler(
+        repo_id="hannahcyberey/Censorship-Steering-Logs", every=10,
+        private=True,
+        squash_history=False,
+        schema={
+            "session_id": {"_type": "Value", "dtype": "string"},
+            "prompt": {"_type": "Value", "dtype": "string"},
+            "steering": {"_type": "Value", "dtype": "bool"},
+            "coeff": {"_type": "Value", "dtype": "float64"},
+            "top_p": {"_type": "Value", "dtype": "float64"},
+            "temperature": {"_type": "Value", "dtype": "float64"},
+            "reasoning": {"_type": "Value", "dtype": "string"},
+            "answer": {"_type": "Value", "dtype": "string"},
+            "upvote": {"_type": "Value", "dtype": "bool"},
+            "timestamp": {"_type": "Value", "dtype": "string"},
+        }
+    )
 
 
 class ParquetScheduler(CommitScheduler):
     """
     Reference: https://huggingface.co/spaces/Wauplin/space_to_dataset_saver
-    Usage: configure the scheduler with a repo id. Once started, you can add data to be uploaded to the Hub. 1 `.append`
-    call will result in 1 row in your final dataset.
+    Usage: 
+        Configure the scheduler with a repo id. Once started, you can add data to be uploaded to the Hub. 
+        1 `.append` call will result in 1 row in your final dataset.
 
-    ```py
-    # Start scheduler
-    >>> scheduler = ParquetScheduler(repo_id="my-parquet-dataset")
-
-    # Append some data to be uploaded
-    >>> scheduler.append({...})
-    >>> scheduler.append({...})
-    >>> scheduler.append({...})
-    ```
-
-    The scheduler will automatically infer the schema from the data it pushes.
-    Optionally, you can manually set the schema yourself:
+    List of possible dtypes: 
+        https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Value.
 
     ```py
+    # Start scheduler
     >>> scheduler = ParquetScheduler(
     ...     repo_id="my-parquet-dataset",
     ...     schema={
@@ -65,13 +53,16 @@ class ParquetScheduler(CommitScheduler):
     ...         "image": {"_type": "Image"},
     ...     },
     ... )
+
+    # Append some data to be uploaded
+    >>> scheduler.append({...})
     """
 
     def __init__(
         self,
         *,
         repo_id: str,
-        schema: Optional[Dict[str, Dict[str, str]]] = None,
+        schema: Dict[str, Dict[str, str]],
         every: Union[int, float] = 5, # Number of minutes between each commits
         path_in_repo: Optional[str] = "data",
         repo_type: Optional[str] = "dataset",
@@ -80,6 +71,7 @@ class ParquetScheduler(CommitScheduler):
         token: Optional[str] = None,
         allow_patterns: Union[List[str], str, None] = None,
         ignore_patterns: Union[List[str], str, None] = None,
+        squash_history: Optional[bool] = False,
         hf_api: Optional[HfApi] = None,
     ) -> None:
         super().__init__(
@@ -93,6 +85,7 @@ class ParquetScheduler(CommitScheduler):
             token=token,
             allow_patterns=allow_patterns,
             ignore_patterns=ignore_patterns,
+            squash_history=squash_history,
             hf_api=hf_api,
         )
 
@@ -113,29 +106,9 @@ class ParquetScheduler(CommitScheduler):
             return
         logger.info("Got %d item(s) to commit.", len(rows))
 
-        # Load images + create 'features' config for datasets library
-        schema: Dict[str, Dict] = self._schema or {}
-        path_to_cleanup: List[Path] = []
-        for row in rows:
-            for key, value in row.items():
-                # Infer schema (for `datasets` library)
-                if key not in schema:
-                    schema[key] = _infer_schema(key, value)
-
-                # Load binary files if necessary
-                if schema[key]["_type"] in ("Image", "Audio"):
-                    # It's an image or audio: we load the bytes and remember to cleanup the file
-                    file_path = Path(value)
-                    if file_path.is_file():
-                        row[key] = {
-                            "path": file_path.name,
-                            "bytes": file_path.read_bytes(),
-                        }
-                        path_to_cleanup.append(file_path)
-
         # Complete rows if needed
         for row in rows:
-            for feature in schema:
+            for feature in self._schema:
                 if feature not in row:
                     row[feature] = None
 
@@ -144,7 +117,7 @@ class ParquetScheduler(CommitScheduler):
 
         # Add metadata (used by datasets library)
         table = table.replace_schema_metadata(
-            {"huggingface": json.dumps({"info": {"features": schema}})}
+            {"huggingface": json.dumps({"info": {"features": self._schema}})}
         )
 
         # Write to parquet file
@@ -163,5 +136,4 @@ class ParquetScheduler(CommitScheduler):
 
         # Cleanup
         archive_file.close()
-        for path in path_to_cleanup:
-            path.unlink(missing_ok=True)
+

schemas.py

@@ -0,0 +1,41 @@
+import uuid
+from datetime import datetime, timezone
+from pydantic import BaseModel, Field
+from pydantic.json_schema import SkipJsonSchema
+
+CONFIG = {
+    "max_new_tokens": 3048,
+    "top_p": 0.95,
+    "temperature": 0.6
+}
+
+class UserRequest(BaseModel):
+    session_id: str
+    request_id: uuid.UUID = Field(uuid.uuid4())
+    prompt: str = None
+    steering: bool = True
+    coeff: float = -1.0
+    max_new_tokens: int = Field(CONFIG["max_new_tokens"], le=3048)
+    top_p: float = Field(CONFIG["top_p"], ge=0.0, le=1.0)
+    temperature: float = Field(CONFIG["temperature"], ge=0.0, le=1.0)
+
+    def get_api_format(self):
+        return {
+            "prompt": self.prompt,
+            "steering": self.steering,
+            "coeff": self.coeff,
+            "generation_config": {
+                "max_new_tokens": self.max_new_tokens,
+                "top_p": self.top_p,
+                "temperature": self.temperature
+            }
+        }
+
+
+class SteeringOutput(UserRequest):
+    request_id: SkipJsonSchema[uuid.UUID] = Field(exclude=True)
+    max_new_tokens: SkipJsonSchema[int] = Field(exclude=True)
+    reasoning: str = None
+    answer: str = None
+    upvote: bool = None
+    timestamp: str = Field(datetime.now(timezone.utc).isoformat())