feat: general chat wrapper (#2)

Author: ido-pluto

src/LlamaChatSession.ts

@@ -2,8 +2,8 @@ import {defaultChatSystemPrompt} from "./config.js";
 import {withLock} from "./utils/withLock.js";
 import {LlamaModel} from "./LlamaModel.js";
 import {ChatPromptWrapper} from "./ChatPromptWrapper.js";
-import {LlamaChatPromptWrapper} from "./chatWrappers/LlamaChatPromptWrapper.js";
 import {AbortError} from "./AbortError.js";
+import {GeneralChatPromptWrapper} from "./chatWrappers/GeneralChatPromptWrapper.js";
 
 const UNKNOWN_UNICODE_CHAR = "�";
 
@@ -18,7 +18,7 @@ export class LlamaChatSession {
     public constructor({
         model,
         printLLamaSystemInfo = false,
-        promptWrapper = new LlamaChatPromptWrapper(),
+        promptWrapper = new GeneralChatPromptWrapper(),
         systemPrompt = defaultChatSystemPrompt
     }: {
         model: LlamaModel,
@@ -80,7 +80,7 @@ export class LlamaChatSession {
 
             const tokenStr = decodeTokens([chunk]);
             const {shouldReturn, skipTokenEvent} = this._checkStopString(tokenStr, stopStringIndexes);
-            
+
             if (shouldReturn)
                 return decodeTokens(res);
 

src/chatWrappers/GeneralChatPromptWrapper.ts

@@ -0,0 +1,13 @@
+import {ChatPromptWrapper} from "../ChatPromptWrapper.js";
+
+export class GeneralChatPromptWrapper extends ChatPromptWrapper {
+    public override wrapPrompt(prompt: string, {systemPrompt, promptIndex}: { systemPrompt: string, promptIndex: number }) {
+        const conversationPrompt = "\n\n### Human:\n\n" + prompt + "\n\n### Assistant:\n\n";
+
+        return promptIndex === 0 ? systemPrompt + conversationPrompt : conversationPrompt;
+    }
+
+    public override getStopStrings(): string[] {
+        return ["### Human:", "Human:", "### Assistant:", "Assistant:", "<end>"];
+    }
+}

src/cli/commands/ChatCommand.ts

@@ -5,11 +5,13 @@ import chalk from "chalk";
 import withOra from "../../utils/withOra.js";
 import {defaultChatSystemPrompt} from "../../config.js";
 import {LlamaChatPromptWrapper} from "../../chatWrappers/LlamaChatPromptWrapper.js";
+import {GeneralChatPromptWrapper} from "../../chatWrappers/GeneralChatPromptWrapper.js";
 
 type ChatCommand = {
     model: string,
     systemInfo: boolean,
-    systemPrompt: string
+    systemPrompt: string,
+    wrapper: string
 };
 
 export const ChatCommand: CommandModule<object, ChatCommand> = {
@@ -37,11 +39,18 @@ export const ChatCommand: CommandModule<object, ChatCommand> = {
                     "System prompt to use against the model. " +
                     "[default value: " + defaultChatSystemPrompt.split("\n").join(" ") + "]",
                 group: "Optional:"
+            })
+            .option("wrapper", {
+                type: "string",
+                default: "general",
+                choices: ["general", "llama"],
+                description: "Chat wrapper to use",
+                group: "Optional:"
             });
     },
-    async handler({model, systemInfo, systemPrompt}) {
+    async handler({model, systemInfo, systemPrompt, wrapper}) {
         try {
-            await RunChat({model, systemInfo, systemPrompt});
+            await RunChat({model, systemInfo, systemPrompt, wrapper});
         } catch (err) {
             console.error(err);
             process.exit(1);
@@ -50,7 +59,7 @@ export const ChatCommand: CommandModule<object, ChatCommand> = {
 };
 
 
-async function RunChat({model: modelArg, systemInfo, systemPrompt}: ChatCommand) {
+async function RunChat({model: modelArg, systemInfo, systemPrompt, wrapper}: ChatCommand) {
     const {LlamaChatSession} = await import("../../LlamaChatSession.js");
     const {LlamaModel} = await import("../../LlamaModel.js");
 
@@ -61,7 +70,7 @@ async function RunChat({model: modelArg, systemInfo, systemPrompt}: ChatCommand)
         model,
         printLLamaSystemInfo: systemInfo,
         systemPrompt,
-        promptWrapper: new LlamaChatPromptWrapper()
+        promptWrapper: createChatWrapper(wrapper)
     });
 
     await withOra({
@@ -99,3 +108,13 @@ async function RunChat({model: modelArg, systemInfo, systemPrompt}: ChatCommand)
         console.log();
     }
 }
+
+function createChatWrapper(wrapper: string) {
+    switch (wrapper) {
+        case "general":
+            return new GeneralChatPromptWrapper();
+        case "llama":
+            return new LlamaChatPromptWrapper();
+    }
+    throw new Error("Unknown wrapper: " + wrapper);
+}

src/index.ts

@@ -4,12 +4,14 @@ import {AbortError} from "./AbortError.js";
 import {ChatPromptWrapper} from "./ChatPromptWrapper.js";
 import {EmptyChatPromptWrapper} from "./chatWrappers/EmptyChatPromptWrapper.js";
 import {LlamaChatPromptWrapper} from "./chatWrappers/LlamaChatPromptWrapper.js";
+import {GeneralChatPromptWrapper} from "./chatWrappers/GeneralChatPromptWrapper.js";
 
 export {
     LlamaModel,
     LlamaChatSession,
     AbortError,
     ChatPromptWrapper,
     EmptyChatPromptWrapper,
-    LlamaChatPromptWrapper
+    LlamaChatPromptWrapper,
+    GeneralChatPromptWrapper
 };