UNPKG

node-llama-cpp

Version:

Run AI models locally on your machine with node.js bindings for llama.cpp. Enforce a JSON schema on the model output on the generation level

95 lines 3.81 kB
import { ChatWrapper } from "../ChatWrapper.js"; import { SpecialToken, LlamaText, SpecialTokensText } from "../utils/LlamaText.js"; // source: https://huggingface.co/blog/llama2#how-to-prompt-llama-2 export class Llama2ChatWrapper extends ChatWrapper { wrapperName = "Llama2Chat"; /** @internal */ _addSpaceBeforeEos; constructor({ addSpaceBeforeEos = false } = {}) { super(); this._addSpaceBeforeEos = addSpaceBeforeEos; } generateContextState({ chatHistory, availableFunctions, documentFunctionParams }) { const historyWithFunctions = this.addAvailableFunctionsSystemMessageToHistory(chatHistory, availableFunctions, { documentParams: documentFunctionParams }); const resultItems = []; let systemTexts = []; let userTexts = []; let modelTexts = []; let currentAggregateFocus = null; function flush() { if (systemTexts.length > 0 || userTexts.length > 0 || modelTexts.length > 0) resultItems.push({ system: LlamaText.joinValues("\n\n", systemTexts), user: LlamaText.joinValues("\n\n", userTexts), model: LlamaText.joinValues("\n\n", modelTexts) }); systemTexts = []; userTexts = []; modelTexts = []; } for (const item of historyWithFunctions) { if (item.type === "system") { if (currentAggregateFocus !== "system") flush(); currentAggregateFocus = "system"; systemTexts.push(LlamaText.fromJSON(item.text)); } else if (item.type === "user") { if (currentAggregateFocus !== "system" && currentAggregateFocus !== "user") flush(); currentAggregateFocus = "user"; userTexts.push(LlamaText(item.text)); } else if (item.type === "model") { currentAggregateFocus = "model"; modelTexts.push(this.generateModelResponseText(item.response)); } else void item; } flush(); const contextText = LlamaText(resultItems.map(({ system, user, model }, index) => { const isLastItem = index === resultItems.length - 1; return LlamaText([ new SpecialToken("BOS"), (system.values.length === 0 && user.values.length === 0) ? LlamaText([]) : LlamaText([ new SpecialTokensText("[INST] "), system.values.length === 0 ? LlamaText([]) : LlamaText([ new SpecialTokensText("<<SYS>>\n"), system, new SpecialTokensText("\n<</SYS>>\n\n") ]), user, new SpecialTokensText(" [/INST] ") ]), model, this._addSpaceBeforeEos ? " " : "", isLastItem ? LlamaText([]) : new SpecialToken("EOS") ]); })); return { contextText, stopGenerationTriggers: [ LlamaText(new SpecialToken("EOS")), LlamaText("</s>") ] }; } /** @internal */ static _getOptionConfigurationsToTestIfCanSupersedeJinjaTemplate() { return [ { addSpaceBeforeEos: false }, { addSpaceBeforeEos: true } ]; } } //# sourceMappingURL=Llama2ChatWrapper.js.map