Format templates, add pwa icon

build/html.ts

@@ -70,7 +70,7 @@ export async function buildHTML(game: string, { production = false, mobile = fal
         }
         let manifest = '';
         if (production && !local) {
-            const pwaIcon = `data:;base64,${await b64(pwaIconFile)}`;
+            const pwaIcon = `data:image/png;base64,${await b64(pwaIconFile)}`;
             const publishURL = process.env.PUBLISH_URL ? `${process.env.PUBLISH_URL}${game}` : '.';
             const manifestJSON = JSON.stringify({
                 name: title,
@@ -87,7 +87,7 @@ export async function buildHTML(game: string, { production = false, mobile = fal
                     type: 'image/png'
                 }]
             });
-            manifest = `<link rel="manifest" href="data:;base64,${await b64(manifestJSON)}" />`;
+            manifest = `<link rel="manifest" href="data:application/json;base64,${await b64(manifestJSON)}" />`;
         }
         let script = await scriptFile.text();
         const inits = new Set<string>();

src/games/ai-story/assets/style.css

@@ -77,6 +77,7 @@ body {
     height: 100dvh;
     font-size: 16px;
     line-height: 1.5;
+    touch-action: none;
 
     .root {
         background-size: cover;

src/games/ai-story/components/header/connectionEditor.tsx

@@ -1,10 +1,10 @@
 import { useCallback, useEffect, useMemo, useState } from 'preact/hooks';
 import styles from './header.module.css';
 import { Connection, HORDE_ANON_KEY, type IConnection, type IHordeModel } from '../../tools/connection';
-import { Instruct } from '../../contexts/state';
 import { useInputState } from '@common/hooks/useInputState';
 import { useInputCallback } from '@common/hooks/useInputCallback';
 import { Huggingface } from '../../tools/huggingface';
+import { INSTRUCTS } from '../../contexts/state';
 
 interface IProps {
     connection: IConnection;
@@ -50,7 +50,6 @@ export const ConnectionEditor = ({ connection, setConnection }: IProps) => {
                 .then(template => {
                     if (template) {
                         setModelTemplate(template);
-                        setInstruct(template);
                     }
                 });
         }
@@ -109,15 +108,18 @@ export const ConnectionEditor = ({ connection, setConnection }: IProps) => {
                     <option value={modelTemplate} title='Native for model'>{modelName}</option>
                 </optgroup>}
                 <optgroup label='Manual templates'>
-                    {Object.entries(Instruct).map(([label, value]) => (
+                    {Object.entries(INSTRUCTS).map(([label, value]) => (
                         <option value={value} key={value}>
-                            {label.toLowerCase()}
+                            {label}
                         </option>
                     ))}
                 </optgroup>
-                {instruct !== modelTemplate && <optgroup label='Custom'>
+                {instruct !== modelTemplate
-                    <option value={connection.instruct}>Custom</option>
+                    && !Object.values(INSTRUCTS).includes(instruct)
-                </optgroup>}
+                    && <optgroup label='Custom'>
+                        <option value={connection.instruct}>Custom</option>
+                    </optgroup>
+                }
             </select>
             {connection.type === 'kobold' && <input
                 value={connectionUrl}

src/games/ai-story/components/header/header.module.css

@@ -27,6 +27,12 @@
         justify-content: center;
         align-items: center;
         gap: 16px;
+
+        .notImportant {
+            @media (width <= 600px) {
+                display: none;
+            }
+        }
     }
 
     .buttons {

src/games/ai-story/components/header/header.tsx

@@ -13,7 +13,7 @@ import { ConnectionEditor } from "./connectionEditor";
 import styles from './header.module.css';
 
 export const Header = () => {
-    const { contextLength, promptTokens, modelName, spentKudos } = useContext(LLMContext);
+    const { contextLength, promptTokens, modelName, spentKudos, hasToolCalls } = useContext(LLMContext);
     const {
         messages,
         connection,
@@ -102,11 +102,11 @@ export const Header = () => {
                     </button>
                 </div>
                 <div class={styles.info}>
-                    <span>{modelName}</span>
+                    <span class={styles.notImportant}>{modelName}</span>
                     <span>📃{promptTokens}/{contextLength}</span>
                     {connection.type === 'horde' ? <>
-                        <span>💲{spentKudos}</span>
+                        <span class={styles.notImportant}>💲{spentKudos}</span>
-                        <span>💰{totalSpentKudos}</span>
+                        <span class={styles.notImportant}>💰{totalSpentKudos}</span>
                     </> : null}
                 </div>
             </div>
@@ -181,7 +181,10 @@ export const Header = () => {
                         &nbsp;Enable summarization
                     </label>
                     <hr />
-                    <h4 class={styles.modalTitle}>Instruct template</h4>
+                    <h4 class={styles.modalTitle}>
+                        Instruct template
+                        {hasToolCalls && <small> (tool calls)</small>}
+                    </h4>
                     <Ace value={connection.instruct} onInput={setInstruct} />
                 </div>
             </Modal>

src/games/ai-story/components/message/message.tsx

@@ -33,7 +33,7 @@ export const Message = ({ message, index, isLastUser, isLastAssistant }: IProps)
     }, [content]);
 
     const handleSaveEdit = useCallback(() => {
-        editMessage(index, editedMessage.trim(), cost);
+        editMessage(index, editedMessage.trim(), 0);
         editSummary(index, '', 0);
         setEditing(false);
     }, [editMessage, editSummary, index, editedMessage, cost]);

src/games/ai-story/contexts/llm.tsx

@@ -227,7 +227,6 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
 
             let messageId = messages.length - 1;
             let text = '';
-            let cost = 0;
 
             const { prompt, isRegen } = await actions.compilePrompt(messages, { continueLast });
 
@@ -242,14 +241,13 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
             editSummary(messageId, 'Generating...', 0);
             for await (const chunk of actions.generate(prompt)) {
                 text += chunk.text;
-                cost += chunk.cost;
                 setPromptTokens(promptTokens + approximateTokens(text));
-                editMessage(messageId, text.trim(), cost);
+                editMessage(messageId, text.trim(), chunk.cost);
             }
             generating.setFalse();
 
             text = MessageTools.trimSentence(text);
-            editMessage(messageId, text, cost);
+            editMessage(messageId, text, 0);
             editSummary(messageId, '', 0);
 
             MessageTools.playReady();

src/games/ai-story/contexts/state.tsx

@@ -6,6 +6,7 @@ import { type IConnection } from "../tools/connection";
 import { loadObject, saveObject } from "../tools/storage";
 import { useInputCallback } from "@common/hooks/useInputCallback";
 import { callUpdater, throttle } from "@common/utils";
+import { Huggingface } from "../tools/huggingface";
 
 interface IStory {
     lore: string;
@@ -96,32 +97,39 @@ interface IActions {
 const SAVE_KEY = 'ai_game_save_state';
 export const DEFAULT_STORY = 'default';
 
-export enum Instruct {
+const INSTRUCT_MISTRAL = Huggingface.formatTemplate(`{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% set loop_messages = messages[1:] %}{% else %}{% set loop_messages = messages %}{% endif %}{% for message in loop_messages %}{% if message['role'] == 'user' %}{% if loop.first and system_message is defined %}{{ ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST]' }}{% else %}{{ ' [INST] ' + message['content'] + ' [/INST]' }}{% endif %}{% elif message['role'] == 'assistant' %}{{ ' ' + message['content'] + '</s>' }}{% endif %}{% endfor %}`);
-    CHATML = `{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n\\n' + message['content'] + '<|im_end|>' + '\\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n\\n' }}{% endif %}`,
 
-    LLAMA = `{% for message in messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n' + message['content'] | trim + '<|eot_id|>' %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}{% endif %}`,
+const INSTRUCT_CHATML = Huggingface.formatTemplate(`{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\n\n' + message['content'] + '<|im_end|>' + '\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n\n' }}{% endif %}`);
 
-    MISTRAL = `{%- if messages[0]['role'] == 'system' %}{%- set system_message = messages[0]['content'] %}{%- set loop_messages = messages[1:] %}{%- else %}{%- set loop_messages = messages %}{%- endif %}{%- for message in loop_messages %}{%- if message['role'] == 'user' %}{%- if loop.first and system_message is defined %}{{- ' [INST] ' + system_message + '\\n\\n' + message['content'] + ' [/INST]' }}{%- else %}{{- ' [INST] ' + message['content'] + ' [/INST]' }}{%- endif %}{%- elif message['role'] == 'assistant' %}{{- ' ' + message['content'] + '</s>'}}{%- endif %}{%- endfor %}`,
+const INSTRUCT_LLAMA = Huggingface.formatTemplate(`{% for message in messages %}{{ '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}`);
 
-    METHARME = `{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{'<|system|>' + message['content'] }}{% elif message['role'] == 'user' %}{{'<|user|>' + message['content'] }}{% elif message['role'] == 'assistant' %}{{'<|model|>' + message['content'] }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|model|>' }}{% endif %}`,
+const INSTRUCT_ALPACA = Huggingface.formatTemplate(`{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{ message['content'] + '\n\n' }}{% elif message['role'] == 'user' %}{{ '### Instruction:\n\n' + message['content'] + '\n\n' }}{% elif message['role'] == 'assistant' %}{{ '### Response:\n\n' + message['content'] + '\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '### Response:\n\n' }}{% endif %}`);
 
-    GEMMA = `{% for message in messages %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model\n'}}{% endif %}`,
+const INSTRUCT_METHARME = Huggingface.formatTemplate(`{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{ '<|system|>' + message['content'] }}{% elif message['role'] == 'user' %}{{ '<|user|>' + message['content'] }}{% elif message['role'] == 'assistant' %}{{'<|model|>' + message['content'] }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|model|>' }}{% endif %}`);
 
-    ALPACA = `{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{ message['content'] + '\\n\\n'}}{% elif message['role'] == 'user' %}{{'### Instruction:\\n\\n' + message['content'] + '\\n\\n'}}{% elif message['role'] == 'assistant' %}{{'### Response:\\n\\n' + message['content'] + '\\n\\n'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '### Response:\\n\\n' }}{% endif %}`,
+const INSTRUCT_GEMMA = Huggingface.formatTemplate(`{% for message in messages %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<start_of_turn>model\n' }}{% endif %}`);
-};
+
+export const INSTRUCTS = {
+    'Mistral': INSTRUCT_MISTRAL,
+    'ChatML': INSTRUCT_CHATML,
+    'LLama': INSTRUCT_LLAMA,
+    'Alpaca': INSTRUCT_ALPACA,
+    'Metharme': INSTRUCT_METHARME,
+    'Gemma': INSTRUCT_GEMMA,
+}
 
 const DEFAULT_CONTEXT: IContext = {
     currentConnection: 0,
     availableConnections: [{
         type: 'kobold',
         url: 'http://localhost:5001',
-        instruct: Instruct.MISTRAL,
+        instruct: INSTRUCT_MISTRAL,
     }],
     input: '',
     systemPrompt: 'You are a creative writer. Write a story based on the world description below. Story should be adult and mature; and could include swearing, violence and unfairness. Portray characters realistically and stay in the lore.',
     stories: {},
     currentStory: DEFAULT_STORY,
-    userPrompt: `{% if isStart -%}
+    userPrompt: Huggingface.formatTemplate(`{% if isStart -%}
 Write a novel using information above as a reference.
 {%- else -%}
 Continue the story forward.
@@ -131,7 +139,7 @@ Continue the story forward.
 This is the description of what should happen next in your answer: {{ prompt | trim }}
 {% endif %}
 Remember that this story should be infinite and go forever.
-Make sure to follow the world description and rules exactly. Avoid cliffhangers and pauses, be creative.`,
+Make sure to follow the world description and rules exactly. Avoid cliffhangers and pauses, be creative.`),
     summarizePrompt: 'Summarize following text in one paragraph:\n\n{{ message }}\n\nAnswer with shortened text only.',
     summaryEnabled: true,
     bannedWords: [],

src/games/ai-story/tools/huggingface.ts

@@ -110,6 +110,20 @@ export namespace Huggingface {
                 });
                 if (fileResponse) {
                     const maybeConfig = JSON.parse(await fileResponse.text());
+                    if (!hasField(maybeConfig, 'chat_template') || !maybeConfig.chat_template) {
+                        console.log(`[huggingface] searching template in '${name}/chat_template.jinja'`);
+                        const templateResponse = await hub.downloadFile({
+                            repo: name,
+                            path: 'chat_template.jinja',
+                        }).catch(() => null);
+
+                        if (templateResponse) {
+                            const template = await templateResponse.text().catch(() => null);
+                            if (template) {
+                                maybeConfig.chat_template = template;
+                            }
+                        }
+                    }
                     if (isTokenizerConfig(maybeConfig)) {
                         tokenizerConfig = maybeConfig;
                         foundName = `${name}/tokenizer_config.json`;
@@ -223,6 +237,110 @@ export namespace Huggingface {
         return text.includes(needle);
     }
 
+    export const minifyTemplate = (input: string, config?: TokenizerConfig) => {
+        let minified = input;
+        do {
+            input = minified;
+            minified = input.replace(/raise_exception\(('[^')]+'|"[^")]+")\)/g, `''`)
+                .replace(/(['"])\s*\+\s*bos_token/gi, `$1`)
+                .replace(/bos_token\s*\+\s*(['"])/gi, `$1`)
+                .replace(/(['"])\s*\+\s*eos_token/gi, `${config?.eos_token?.replace('$', '$$') ?? ''}$1`)
+                .replace(/eos_token\s*\+\s*(['"])/gi, `$1${config?.eos_token?.replace('$', '$$') ?? ''}`)
+                .replace(/\{#-?[^#]+-?#}/gi, '')
+                .replace(/\s*(\{[{%])-/gi, '$1')
+                .replace(/-([}%]\})\s*/gi, '$1')
+                .replace(/\{\{\s*(''|"")\s*\}\}/g, '')
+                .replace(/\s*\}\}\{\{\s*/, ' + ')
+                .replace(/\n+['"]/g, (match) => match.replace(/\n/gi, '\\n'))
+                .replace(/'\s*\+\s*'/g, '')
+                .replace(/"\s*\+\s*"/g, '')
+                .replace(/\{%\s*else\s*%\}\{%\s*endif\s*%\}/gi, '{% endif %}')
+                .replace(/\{%\s*elif[^}]+%\}\{%\s*endif\s*%\}/gi, '{% endif %}')
+                .replace(/\{%\s*if[^}]+%\}\{%\s*endif\s*%\}/gi, '')
+                .replaceAll('bos_token', `''`)
+                .replaceAll('eos_token', `'${config?.eos_token ?? ''}'`);
+        } while (minified !== input);
+
+        return minified;
+    }
+
+    export const formatTemplate = (input: string, config?: TokenizerConfig) => {
+        const minified = minifyTemplate(input, config);
+
+        type ParserState = 'none' | 'open_brace' | 'block' | 'block_end' | 'quote' | 'escaped';
+        let state: ParserState = 'none';
+        let currentBlock = '';
+        let blockStart = '';
+        let quoteStart = '';
+        let escaped = false;
+
+        const blocks: string[] = [];
+
+        for (const ch of minified) {
+            currentBlock += ch;
+            if (state === 'none') {
+                if (ch === '{') {
+                    state = 'open_brace';
+                }
+            } else if (state === 'open_brace') {
+                if (ch === '{' || ch === '%') {
+                    blockStart = ch;
+                    state = 'block';
+                    currentBlock += '-';
+                } else {
+                    state = 'none';
+                }
+            } else if (state === 'block') {
+                if (ch === '"' || ch === "'") {
+                    quoteStart = ch;
+                    state = 'quote';
+                } else if (ch === blockStart || blockStart === '{' && ch === '}') {
+                    currentBlock = currentBlock.slice(0, -1) + '-' + ch;
+                    state = 'block_end';
+                }
+            } else if (state === 'block_end') {
+                if (ch === '}') {
+                    state = 'none';
+                    blocks.push(currentBlock);
+                    currentBlock = '';
+                } else {
+                    state = 'block';
+                }
+            } else if (state === 'quote') {
+                if (!escaped && ch === quoteStart) {
+                    state = 'block';
+                } else if (!escaped && ch === '\\') {
+                    escaped = true;
+                } else {
+                    escaped = false;
+                }
+            }
+        }
+        if (currentBlock) {
+            blocks.push(currentBlock);
+        }
+        let indent = '';
+        for (let i = 0; i < blocks.length; i++) {
+            const line = blocks[i];
+            const content = line.slice(3).trim();
+            if (content.startsWith('if ') || content.startsWith('for ')) {
+                blocks[i] = indent + line;
+                indent += '  ';
+            } else if (content.startsWith('else ') || content.startsWith('elif ')) {
+                indent = indent.slice(2);
+                blocks[i] = indent + line;
+                indent += '  ';
+            } else if (content.startsWith("end")) {
+                indent = indent.slice(2);
+                blocks[i] = indent + line;
+            } else {
+                blocks[i] = indent + line;
+            }
+        }
+
+        return blocks.filter(b => b.trim()).join('\n');
+    }
+
     export const findModelTemplate = async (modelName: string): Promise<string | null> => {
         modelName = normalizeModel(modelName);
         if (!modelName) return '';
@@ -235,18 +353,7 @@ export namespace Huggingface {
             const config = await loadHuggingfaceTokenizerConfig(modelName);
 
             if (config?.chat_template?.trim()) {
-                template = config.chat_template.trim()
+                template = formatTemplate(config.chat_template, config);
-                    .replace(/raise_exception\(('[^')]+'|"[^")]+")\)/g, `''`)
-                    .replaceAll('eos_token', `'${config.eos_token ?? ''}'`)
-                    .replaceAll('bos_token', `''`)
-                    .replace(/\{\{ ?(''|"") ?\}\}/g, '')
-                    .replace(/\n'/g, `\\n'`)
-                    .replace(/\n"/g, `\\n"`)
-                    .replace(/'\s*\+\s*'/g, '')
-                    .replace(/"\s*\+\s*"/g, '')
-                    .replace(/\{%\s*else\s*%\}\{%\s*endif\s*%\}/gi, '{% endif %}')
-                    .replace(/\{%\s*elif[^}]+%\}\{%\s*endif\s*%\}/gi, '{% endif %}')
-                    .replace(/\{%\s*if[^}]+%\}\{%\s*endif\s*%\}/gi, '');
             }
         }
 

src/games/ai-story/tools/messages.ts

@@ -103,7 +103,7 @@ export namespace MessageTools {
             (m, i) => ({
                 ...m,
                 swipes: i === index
-                    ? m.swipes.map((s, si) => (si === m.currentSwipe ? { ...s, ...update, cost: s.cost + cost } : s))
+                    ? m.swipes.map((s, si) => (si === m.currentSwipe ? { ...s, ...update, cost: (s.cost || 0) + cost } : s))
                     : m.swipes
             })
         )