85 lines
3.5 KiB
JavaScript
85 lines
3.5 KiB
JavaScript
import { ChatWrapper } from "../ChatWrapper.js";
|
|
import { SpecialToken, LlamaText, SpecialTokensText } from "../utils/LlamaText.js";
|
|
// source: https://github.com/openai/openai-python/blob/120d225b91a8453e15240a49fb1c6794d8119326/chatml.md
|
|
export class ChatMLChatWrapper extends ChatWrapper {
|
|
wrapperName = "ChatML";
|
|
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") {
|
|
flush();
|
|
currentAggregateFocus = null;
|
|
userTexts.push(LlamaText(item.text));
|
|
}
|
|
else if (item.type === "model") {
|
|
flush();
|
|
currentAggregateFocus = null;
|
|
modelTexts.push(this.generateModelResponseText(item.response));
|
|
}
|
|
else
|
|
void item;
|
|
}
|
|
flush();
|
|
const contextText = LlamaText(new SpecialToken("BOS"), resultItems.map(({ system, user, model }, index) => {
|
|
const isLastItem = index === resultItems.length - 1;
|
|
return LlamaText([
|
|
(system.values.length === 0)
|
|
? LlamaText([])
|
|
: LlamaText([
|
|
new SpecialTokensText("<|im_start|>system\n"),
|
|
system,
|
|
new SpecialTokensText("<|im_end|>\n")
|
|
]),
|
|
(user.values.length === 0)
|
|
? LlamaText([])
|
|
: LlamaText([
|
|
new SpecialTokensText("<|im_start|>user\n"),
|
|
user,
|
|
new SpecialTokensText("<|im_end|>\n")
|
|
]),
|
|
(model.values.length === 0 && !isLastItem)
|
|
? LlamaText([])
|
|
: LlamaText([
|
|
new SpecialTokensText("<|im_start|>assistant\n"),
|
|
model,
|
|
isLastItem
|
|
? LlamaText([])
|
|
: new SpecialTokensText("<|im_end|>\n")
|
|
])
|
|
]);
|
|
}));
|
|
return {
|
|
contextText,
|
|
stopGenerationTriggers: [
|
|
LlamaText(new SpecialToken("EOS")),
|
|
LlamaText(new SpecialTokensText("<|im_end|>")),
|
|
LlamaText("<|im_end|>")
|
|
]
|
|
};
|
|
}
|
|
}
|
|
//# sourceMappingURL=ChatMLChatWrapper.js.map
|