Refactor (9)

index.ts

@@ -2,7 +2,13 @@ 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 { LlmSession, createGrammar, getModel, parseResponse } from "./lib/llm";
+import {
+  LlmSession,
+  createBias,
+  createGrammar,
+  getModel,
+  parseResponse,
+} from "./lib/llm";
 import {
   expandReplyTree,
   getNotes,
@@ -30,9 +36,11 @@ const modelName =
 console.log(`* loading model '${modelName}'`);
 const model = await getModel(modelName);
 const grammar = await createGrammar("あるびのちゃん");
+const bias = createBias(model);
 const baseChatPromptOptions = {
   grammar,
   maxTokens: 256,
+  tokenBias: bias,
 } as const satisfies LLamaChatPromptOptions;
 
 const getSystemPrompt = (
@@ -58,7 +66,7 @@ async function rephrase(text: string) {
   await using rephraseSession = new LlmSession(
     model,
     getSystemPrompt(
-      "user が与えたテキストを『ですます調』（丁寧な文体）で言い換えたものを、そのまま出力してください。",
+      "ユーザが与えたテキストを「～です」「～ます」調（丁寧な文体）で言い換えたものを、そのまま出力してください。",
     ),
   );
   await rephraseSession.init();

lib/llm.ts

@@ -7,6 +7,7 @@ import {
   type LLamaChatPromptOptions,
   LlamaChatSession,
   type LlamaModel,
+  TokenBias,
   createModelDownloader,
   getLlama,
   resolveChatWrapper,
@@ -27,6 +28,15 @@ export async function getModel(model: string) {
   return await llama.loadModel({ modelPath });
 }
 
+export function createBias(model: LlamaModel) {
+  const customBias = new TokenBias(model.tokenizer);
+  for (const token of model.iterateAllTokens()) {
+    const text = model.detokenize([token]);
+    if (text === "{") customBias.set(token, -0.9); // suppress JSON string
+  }
+  return customBias;
+}
+
 export const createGrammar = (assistantName: string) =>
   llama.createGrammarForJsonSchema({
     type: "object",