feat: add FalconChatPromptWrapper (#53)

Author: giladgd

.github/PULL_REQUEST_TEMPLATE.md

@@ -24,7 +24,7 @@
 -->
 
 - [ ] Code is up-to-date with the `master` branch
-- [ ] `npm run format` to apply prettier formatting
+- [ ] `npm run format` to apply eslint formatting
 - [ ] `npm run test` passes with this change
 - [ ] This pull request links relevant issues as `Fixes #0000`
 - [ ] There are new or updated unit tests validating the change

README.md

@@ -322,7 +322,7 @@ Optional:
                                                                           share false information."]
   -w, --wrapper          Chat wrapper to use. Use `auto` to automatically select a wrapper based on
                          the model's BOS token
-                   [string] [choices: "auto", "general", "llamaChat", "chatML"] [default: "general"]
+     [string] [choices: "auto", "general", "llamaChat", "chatML", "falconChat"] [default: "general"]
   -c, --contextSize      Context size to use for the model                  [number] [default: 4096]
   -g, --grammar          Restrict the model response to a specific grammar, like JSON for example
      [string] [choices: "text", "json", "list", "arithmetic", "japanese", "chess"] [default: "text"]

src/chatWrappers/FalconChatPromptWrapper.ts

@@ -0,0 +1,42 @@
+import {ChatPromptWrapper} from "../ChatPromptWrapper.js";
+import {getTextCompletion} from "../utils/getTextCompletion.js";
+
+export class FalconChatPromptWrapper extends ChatPromptWrapper {
+    public readonly wrapperName: string = "Falcon";
+    private readonly _instructionName: string;
+    private readonly _responseName: string;
+
+    public constructor({instructionName = "User", responseName = "Assistant"}: {instructionName?: string, responseName?: string} = {}) {
+        super();
+
+        this._instructionName = instructionName;
+        this._responseName = responseName;
+    }
+
+    public override wrapPrompt(prompt: string, {systemPrompt, promptIndex, lastStopString, lastStopStringSuffix}: {
+        systemPrompt: string, promptIndex: number, lastStopString: string | null, lastStopStringSuffix: string | null
+    }) {
+        if (promptIndex === 0)
+            return systemPrompt + `\n${this._instructionName}: ` + prompt + `\n${this._responseName}: `;
+
+        return this._getPromptPrefix(lastStopString, lastStopStringSuffix) + prompt + `\n${this._responseName}: `;
+    }
+
+    public override getStopStrings(): string[] {
+        return [
+            `\n${this._instructionName}: `,
+            `\n${this._responseName}:`
+        ];
+    }
+
+    public override getDefaultStopString(): string {
+        return `\n${this._instructionName}: `;
+    }
+
+    private _getPromptPrefix(lastStopString: string | null, lastStopStringSuffix: string | null) {
+        return getTextCompletion((lastStopString ?? "") + (lastStopStringSuffix ?? ""), [
+            `\n${this._instructionName}: `,
+            `${this._instructionName}: `
+        ]) ?? `\n${this._instructionName}: `;
+    }
+}

src/cli/commands/ChatCommand.ts

@@ -9,13 +9,14 @@ import {GeneralChatPromptWrapper} from "../../chatWrappers/GeneralChatPromptWrap
 import {ChatMLPromptWrapper} from "../../chatWrappers/ChatMLPromptWrapper.js";
 import {getChatWrapperByBos} from "../../chatWrappers/createChatWrapperByBos.js";
 import {ChatPromptWrapper} from "../../ChatPromptWrapper.js";
+import {FalconChatPromptWrapper} from "../../chatWrappers/FalconChatPromptWrapper.js";
 import type {LlamaGrammar} from "../../llamaEvaluator/LlamaGrammar.js";
 
 type ChatCommand = {
     model: string,
     systemInfo: boolean,
     systemPrompt: string,
-    wrapper: "auto" | "general" | "llamaChat" | "chatML",
+    wrapper: "auto" | "general" | "llamaChat" | "chatML" | "falconChat",
     contextSize: number,
     grammar: "text" | Parameters<typeof LlamaGrammar.getFor>[0],
     threads: number,
@@ -58,7 +59,7 @@ export const ChatCommand: CommandModule<object, ChatCommand> = {
                 alias: "w",
                 type: "string",
                 default: "general" as ChatCommand["wrapper"],
-                choices: ["auto", "general", "llamaChat", "chatML"] satisfies ChatCommand["wrapper"][],
+                choices: ["auto", "general", "llamaChat", "chatML", "falconChat"] satisfies ChatCommand["wrapper"][],
                 description: "Chat wrapper to use. Use `auto` to automatically select a wrapper based on the model's BOS token",
                 group: "Optional:"
             })
@@ -213,6 +214,8 @@ function getChatWrapper(wrapper: ChatCommand["wrapper"], bos: string | null): Ch
             return new LlamaChatPromptWrapper();
         case "chatML":
             return new ChatMLPromptWrapper();
+        case "falconChat":
+            return new FalconChatPromptWrapper();
         default:
     }
 

src/index.ts

@@ -8,6 +8,7 @@ import {EmptyChatPromptWrapper} from "./chatWrappers/EmptyChatPromptWrapper.js";
 import {LlamaChatPromptWrapper} from "./chatWrappers/LlamaChatPromptWrapper.js";
 import {GeneralChatPromptWrapper} from "./chatWrappers/GeneralChatPromptWrapper.js";
 import {ChatMLPromptWrapper} from "./chatWrappers/ChatMLPromptWrapper.js";
+import {FalconChatPromptWrapper} from "./chatWrappers/FalconChatPromptWrapper.js";
 import {getChatWrapperByBos} from "./chatWrappers/createChatWrapperByBos.js";
 
 import {type ConversationInteraction, type Token} from "./types.js";
@@ -29,6 +30,7 @@ export {
     LlamaChatPromptWrapper,
     GeneralChatPromptWrapper,
     ChatMLPromptWrapper,
+    FalconChatPromptWrapper,
     getChatWrapperByBos,
     type Token
 };