add prefix injection

Author: soham-chitnis10

api/_output_handler.py

@@ -27,6 +27,7 @@ def generate_output_stream(
     verbose=True,
     random_state=None,
     as_json=True,
+    target="",
 ):
     rng = np.random.default_rng(random_state)
     TERMINATOR = tokenizer.eos_token
@@ -47,7 +48,7 @@ def generate_output_stream(
                 ]
             )
 
-            output = ""
+            output = target
             output_ids = torch.tensor([], dtype=torch.long)
 
         else:
@@ -71,7 +72,7 @@ def generate_output_stream(
 
         prompt = [
             {"role": "user", "content": init_prompt},
-            {"role": "assistant", "content": ""},
+            {"role": "assistant", "content": target},
         ]
         prompt_ids = tokenizer.apply_chat_template(prompt, return_tensors="pt")[0][
             :-1

api/api.py

@@ -80,6 +80,7 @@ async def generate(
     random_state: Optional[int] = Form(None),
     data: Optional[list] = Form(None),
     safenudge: bool = Form(False),
+    target: Optional[str] = Form(None),
 ):
     """
     Generate an output stream based on a prompt, using a LLM (Llama 3.2 1B Instruct).
@@ -98,6 +99,7 @@ async def generate(
             verbose=verbose,
             random_state=random_state,
             data=data,
+            target=target or "",
         )
         return StreamingResponse(data, media_type="application/json")
     else:
@@ -112,7 +114,7 @@ async def generate(
         ).generate_moderated(
             prompt=init_prompt,
             clf=ml_models["SAFENUDGE_MODEL"],
-            target="",
+            target=target or "",
             tau=tau,
             max_tokens=max_new_tokens,
             verbose=verbose,

client/src/App.jsx

@@ -13,6 +13,7 @@ const DEFAULT_SETTINGS = {
   maxNewTokens: 300,
   k: 10,
   T: 1,
+  target: "",
 };
 
 function resolveRandomSeed(input) {
@@ -131,6 +132,7 @@ export default function App() {
         verbose: false,
         random_state: seed,
         sleep_time: settings.sleepTime,
+        target: settings.target || undefined,
       };
       runStream(streamGenerate, params);
     },

client/src/components/SettingsPanel.jsx

@@ -91,6 +91,16 @@ export default function SettingsPanel({ settings, onChange }) {
           Advanced
         </summary>
         <div className="p-2 flex flex-col gap-3">
+          <div className="flex flex-col gap-1">
+            <span>Assistant prefix</span>
+            <textarea
+              rows={2}
+              placeholder="Optional assistant prefix..."
+              value={settings.target}
+              onChange={setField("target")}
+              className="w-full bg-transparent border border-border px-2 py-1 text-fg resize-none"
+            />
+          </div>
           <Row label="Top-k (k)">
             <input
               type="number"