Refactor

index.ts

@@ -2,15 +2,22 @@ import { parseArgs } from "node:util";
 import { Stream } from "misskey-js";
 import type { Note } from "misskey-js/entities.js";
 import type { ChatHistoryItem, LLamaChatPromptOptions } from "node-llama-cpp";
-import { createGrammar, getModel, LlmSession, parseResponse } from "./lib/llm";
+import { createHarness, type HarnessName, isHarnessName } from "./lib/harness";
+import { createGrammar, getModel, LlmSession } from "./lib/llm";
+import { expandReplyTree, getNotes, me, misskey } from "./lib/misskey";
 import {
-  expandReplyTree,
-  getNotes,
-  me,
-  misskey,
-  sanitizeText,
-} from "./lib/misskey";
-import { sleep } from "./lib/util";
+  formatConversationNote,
+  formatTimeline,
+  postJobPrompt,
+  replyJobPrompt,
+} from "./lib/prompt";
+import { ensurePolite, isPolitenessMode } from "./lib/rephrase";
+import {
+  getSamplingOptions,
+  isSamplingProfile,
+  type SamplingProfile,
+} from "./lib/sampling";
+import { sanitizeText, sleep } from "./lib/util";
 
 const { values } = parseArgs({
   args: Bun.argv,
@@ -20,88 +27,48 @@ const { values } = parseArgs({
       short: "t",
       default: false,
     },
+    harness: {
+      type: "string",
+      default: Bun.env["HARNESS"] ?? "tagged",
+    },
+    politeness: {
+      type: "string",
+      default: Bun.env["POLITENESS"] ?? "auto",
+    },
+    sampling: {
+      type: "string",
+      default: Bun.env["SAMPLING"] ?? "creative",
+    },
   },
   strict: true,
   allowPositionals: true,
 });
 
+if (!isHarnessName(values.harness)) {
+  throw new Error(`Unknown harness: ${values.harness}`);
+}
+const harnessName: HarnessName = values.harness;
+if (!isPolitenessMode(values.politeness)) {
+  throw new Error(`Unknown politeness mode: ${values.politeness}`);
+}
+const politenessMode = values.politeness;
+if (!isSamplingProfile(values.sampling)) {
+  throw new Error(`Unknown sampling profile: ${values.sampling}`);
+}
+const samplingProfile: SamplingProfile = values.sampling;
+
 const modelName =
   Bun.env["MODEL"] ?? "mradermacher/gemma-2-baku-2b-it-GGUF:IQ4_XS";
 console.log(`* loading model '${modelName}'`);
 const model = await getModel(modelName);
 const grammar = await createGrammar("あるびのちゃん");
+const harness = createHarness(harnessName, grammar);
+console.log(`* output harness: ${harness.name}`);
+console.log(`* sampling profile: ${samplingProfile}`);
 const baseChatPromptOptions = {
-  grammar,
-  maxTokens: 256,
-  repeatPenalty: {
-    punishTokensFilter: (tokens) =>
-      tokens.filter((token) => {
-        const text = model.detokenize([token]).trim();
-        for (const s of ["{", "}", "「", "」"]) {
-          if (text.startsWith(s) || text.endsWith(s)) return true;
-        }
-        return false;
-      }),
-    penalty: 1.2,
-    presencePenalty: 0.5,
-    frequencyPenalty: 0.5,
-  },
+  maxTokens: 192,
 } as const satisfies LLamaChatPromptOptions;
 
-const getSystemPrompt = (
-  instruction: string,
-) => `あなたは人造天使「あるびのちゃん」です。機械的・非人間的に観察します。
-キャラ設定:
-- アルギリア製の白髪赤目な人造天使。非人間的な考え方や思想を持つ。
-- 現在は地球の衛星軌道上からインターネットを観察している。
-出力規則:
-- 「～です」「～ます」などの語尾を使って **丁寧に** 話す。
-- 3センテンス程度で出力する。
-- 出力規則の内容について言及しない。
-
-${instruction}`;
-const postJobPrompt = getSystemPrompt(
-  "以下は SNS のタイムラインです。**タイムラインの話題に言及しつつ**、あるびのちゃんとして何かツイートしてください。",
-);
-const replyJobPrompt = getSystemPrompt(
-  "ユーザがあなたへのメッセージを送ってきています。あるびのちゃんとして、発言に返信してください。",
-);
-
-async function rephrase(text: string) {
-  if (
-    text.includes("です") ||
-    text.includes("ます") ||
-    text.includes("でし") ||
-    text.includes("まし") ||
-    text.includes("ません")
-  ) {
-    return text;
-  }
-  await using rephraseSession = new LlmSession(
-    model,
-    "ユーザが与えたテキストを「～です」「～ます」調（丁寧な文体）で言い換えたものを、そのまま出力してください。",
-  );
-  await rephraseSession.init();
-  const res = parseResponse(
-    grammar,
-    await rephraseSession.prompt(JSON.stringify({ text }), {
-      ...baseChatPromptOptions,
-      customStopTriggers: ["ですます"],
-    }),
-  );
-  return res ?? text;
-}
-
-const formatNote = (n: Note) => {
-  if (n.userId === me.id) {
-    return JSON.stringify({ name: "あるびのちゃん", text: n.text });
-  }
-  return JSON.stringify({
-    name: n.user.name ?? n.user.username,
-    text: n.text,
-  });
-};
-
 type Job =
   // read posts and post a note
   | { type: "post" }
@@ -116,29 +83,32 @@ type Job =
 
 async function processPostJob() {
   const notes = await getNotes(10, 0, 5);
-  const input = notes.map(formatNote).join("\n");
+  const input = formatTimeline(notes, me.id);
   const text = await (async () => {
-    await using postJobSession = new LlmSession(model, postJobPrompt);
-    await postJobSession.init();
-    return await parseResponse(
-      grammar,
-      await postJobSession.prompt(input, {
-        ...baseChatPromptOptions,
-        temperature: 1.25,
-        minP: 0.05,
-        repeatPenalty: {
-          lastTokens: 128,
-          penalty: 1.15,
-        },
-      }),
+    await using postJobSession = new LlmSession(
+      model,
+      postJobPrompt + harness.promptSuffix,
     );
+    await postJobSession.init();
+    const raw = await postJobSession.prompt(input, {
+      ...baseChatPromptOptions,
+      ...harness.options,
+      ...getSamplingOptions(modelName, samplingProfile, "post"),
+    });
+    process.stderr.write("\n");
+    return harness.parse(raw);
   })();
   if (text) {
-    const rephrased = await rephrase(text);
-    if (values.test) return;
+    const output = sanitizeText(
+      (await ensurePolite(model, text, politenessMode)).text,
+    );
+    if (values.test) {
+      console.log(`\n${output}\n`);
+      return;
+    }
     await misskey.request("notes/create", {
       visibility: "public",
-      text: sanitizeText(rephrased),
+      text: output,
     });
   }
 }
@@ -148,37 +118,41 @@ async function processReplyJob(job: Extract<Job, { type: "reply" }>) {
     if (n.userId === me.id) {
       return {
         type: "model",
-        response: [formatNote(n)],
+        response: [harness.formatHistory(n.text ?? "")],
       } as const;
     }
     return {
       type: "user",
-      text: formatNote(n),
+      text: formatConversationNote(n, me.id),
     } as const;
   });
   const text = await (async () => {
-    await using session = new LlmSession(model, replyJobPrompt, history);
-    await session.init();
-    return parseResponse(
-      grammar,
-      await session.prompt(formatNote(job.last), {
-        ...baseChatPromptOptions,
-        temperature: 0.8,
-        minP: 0.1,
-        repeatPenalty: {
-          lastTokens: 128,
-          penalty: 1.15,
-        },
-      }),
+    await using session = new LlmSession(
+      model,
+      replyJobPrompt + harness.promptSuffix,
+      history,
     );
+    await session.init();
+    const raw = await session.prompt(formatConversationNote(job.last, me.id), {
+      ...baseChatPromptOptions,
+      ...harness.options,
+      ...getSamplingOptions(modelName, samplingProfile, "reply"),
+    });
+    process.stderr.write("\n");
+    return harness.parse(raw);
   })();
 
   if (text) {
-    const rephrased = await rephrase(text);
-    if (values.test) return;
+    const output = sanitizeText(
+      (await ensurePolite(model, text, politenessMode)).text,
+    );
+    if (values.test) {
+      console.log(`\n${output}\n`);
+      return;
+    }
     await misskey.request("notes/create", {
       visibility: job.visibility,
-      text: sanitizeText(rephrased),
+      text: output,
       replyId: job.id,
     });
   }
@@ -256,12 +230,12 @@ function initializeStream() {
   });
 }
 
-/** pop from the job queue and run it */
+/** take jobs from the queue in arrival order */
 async function runJob() {
   while (true) {
-    const job = jobs.pop();
+    const job = jobs.shift();
     if (job) {
-      console.log(`* pop: ${job.type}`);
+      console.log(`* take: ${job.type}`);
       try {
         await processJob(job);
         console.log("* job complete");