diff --git a/index.ts b/index.ts index 0770319..d1a4f3d 100644 --- a/index.ts +++ b/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(); diff --git a/lib/llm.ts b/lib/llm.ts index 7ac0b71..37ae2b5 100644 --- a/lib/llm.ts +++ b/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",