448 lines
17 KiB
TypeScript
448 lines
17 KiB
TypeScript
import { gguf } from '@huggingface/gguf';
|
|
import * as hub from '@huggingface/hub';
|
|
import { Template } from '@huggingface/jinja';
|
|
import { Tokenizer } from '@huggingface/tokenizers';
|
|
import { normalizeModel } from './model';
|
|
import { loadObject, saveObject } from './storage';
|
|
|
|
export namespace Huggingface {
|
|
export interface ITemplateMessage {
|
|
role: 'user' | 'assistant' | 'system';
|
|
content: string;
|
|
}
|
|
|
|
interface INumberParameter {
|
|
type: 'number';
|
|
enum?: number[];
|
|
description?: string;
|
|
}
|
|
|
|
interface IStringParameter {
|
|
type: 'string';
|
|
enum?: string[];
|
|
description?: string;
|
|
}
|
|
|
|
interface IArrayParameter {
|
|
type: 'array';
|
|
description?: string;
|
|
items: IParameter;
|
|
}
|
|
|
|
interface IObjectParameter {
|
|
type: 'object';
|
|
description?: string;
|
|
properties: Record<string, IParameter>;
|
|
required?: string[];
|
|
}
|
|
|
|
type IParameter = INumberParameter | IStringParameter | IArrayParameter | IObjectParameter;
|
|
|
|
interface ITool {
|
|
type: 'function',
|
|
function: {
|
|
name: string;
|
|
description?: string;
|
|
parameters?: IObjectParameter;
|
|
}
|
|
}
|
|
|
|
export interface IFunction {
|
|
name: string;
|
|
description?: string;
|
|
parameters?: Record<string, IParameter>;
|
|
}
|
|
|
|
interface TokenizerConfig {
|
|
chat_template: string;
|
|
bos_token?: string;
|
|
eos_token?: string;
|
|
}
|
|
|
|
type TokenizerJson = any;
|
|
type TokenizerInfo = [TokenizerConfig | null, TokenizerJson | null];
|
|
|
|
const TEMPLATE_CACHE_KEY = 'ai_game_template_cache';
|
|
const TOKENIZER_CACHE_KEY = 'ai_game_tokenizer_cache';
|
|
|
|
const templateCache: Record<string, string> = {};
|
|
|
|
const tokenizerCache: Record<string, TokenizerInfo> = {};
|
|
const prevLoading: Promise<unknown> = Promise.all([
|
|
loadObject(TEMPLATE_CACHE_KEY, {}).then(c => Object.assign(templateCache, c)),
|
|
loadObject(TOKENIZER_CACHE_KEY, {}).then(c => Object.assign(tokenizerCache, c)),
|
|
]);
|
|
|
|
const compiledTemplates = new Map<string, Template>();
|
|
const compiledTokenizers = new Map<string, Tokenizer | null>();
|
|
|
|
const hasField = <T extends string>(obj: unknown, field: T): obj is Record<T, unknown> => (
|
|
obj != null && typeof obj === 'object' && (field in obj)
|
|
);
|
|
const isTokenizerConfig = (obj: unknown): obj is TokenizerConfig => (
|
|
hasField(obj, 'chat_template') && (typeof obj.chat_template === 'string')
|
|
&& (!hasField(obj, 'eos_token') || !obj.eos_token || typeof obj.eos_token === 'string')
|
|
&& (!hasField(obj, 'bos_token') || !obj.bos_token || typeof obj.bos_token === 'string')
|
|
);
|
|
|
|
const loadHuggingfaceTokenizer = async (modelName: string, configOnly = false): Promise<TokenizerInfo> => {
|
|
await prevLoading;
|
|
|
|
modelName = normalizeModel(modelName);
|
|
console.log(`[huggingface] searching config for '${modelName}'`);
|
|
|
|
const cachedConfig = tokenizerCache[modelName];
|
|
if (cachedConfig && cachedConfig[0] != null && cachedConfig[1] != null) {
|
|
console.log(`[huggingface] found cached config for '${modelName}'`);
|
|
return cachedConfig;
|
|
}
|
|
|
|
const hubModels = await Array.fromAsync(hub.listModels({ search: { query: modelName }, additionalFields: ['config'] }));
|
|
const models = hubModels.filter(m => {
|
|
if (m.gated) return false;
|
|
if (!normalizeModel(m.name).includes(modelName)) return false;
|
|
|
|
return true;
|
|
}).sort((a, b) => b.downloads - a.downloads);
|
|
|
|
let tokenizerConfig: TokenizerConfig | null = null;
|
|
let tokenizerJson: TokenizerJson | null = null;
|
|
|
|
for (const model of models) {
|
|
const { config, name } = model;
|
|
|
|
if (name.toLowerCase().includes('gguf')) continue;
|
|
|
|
if (!tokenizerJson && !configOnly) {
|
|
try {
|
|
console.log(`[huggingface] searching tokenizer in '${name}/tokenizer.json'`);
|
|
const fileResponse = await hub.downloadFile({
|
|
repo: name,
|
|
path: 'tokenizer.json',
|
|
});
|
|
if (fileResponse) {
|
|
tokenizerJson = JSON.parse(await fileResponse.text());
|
|
console.log(`[huggingface] found tokenizer in '${name}/tokenizer.json'`);
|
|
}
|
|
} catch { }
|
|
}
|
|
|
|
if (!tokenizerConfig) {
|
|
if (hasField(config, 'tokenizer_config') && isTokenizerConfig(config.tokenizer_config)) {
|
|
tokenizerConfig = config.tokenizer_config;
|
|
console.log(`[huggingface] found config for '${modelName}' in '${name}'`);
|
|
}
|
|
}
|
|
|
|
if (!tokenizerConfig) {
|
|
try {
|
|
console.log(`[huggingface] searching config in '${name}/tokenizer_config.json'`);
|
|
const fileResponse = await hub.downloadFile({
|
|
repo: name,
|
|
path: 'tokenizer_config.json',
|
|
});
|
|
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;
|
|
console.log(`[huggingface] found config for '${modelName}' in '${name}/tokenizer_config.json'`);
|
|
break;
|
|
}
|
|
}
|
|
} catch { }
|
|
}
|
|
}
|
|
|
|
if (!tokenizerConfig) {
|
|
for (const model of models) {
|
|
try {
|
|
for await (const file of hub.listFiles({ repo: model.name, recursive: true })) {
|
|
if (file.type !== 'file' || !file.path.endsWith('.gguf')) continue;
|
|
try {
|
|
console.log(`[huggingface] searching config in '${model.name}/${file.path}'`);
|
|
const fileInfo = await hub.fileDownloadInfo({ repo: model.name, path: file.path });
|
|
if (fileInfo?.url) {
|
|
const { metadata } = await gguf(fileInfo.url);
|
|
if ('tokenizer.chat_template' in metadata) {
|
|
const chat_template = metadata['tokenizer.chat_template'];
|
|
const tokens = metadata['tokenizer.ggml.tokens'];
|
|
const bos_token = tokens[metadata['tokenizer.ggml.bos_token_id']];
|
|
const eos_token = tokens[metadata['tokenizer.ggml.eos_token_id']];
|
|
|
|
const maybeConfig = {
|
|
chat_template,
|
|
bos_token,
|
|
eos_token,
|
|
}
|
|
|
|
if (isTokenizerConfig(maybeConfig)) {
|
|
tokenizerConfig = maybeConfig;
|
|
console.log(`[huggingface] found config for '${modelName}' in '${model.name}/${file.path}'`);
|
|
break;
|
|
}
|
|
} else if ('tokenizer.ggml.model' in metadata) {
|
|
break; // no reason to touch different quants
|
|
}
|
|
}
|
|
} catch { }
|
|
}
|
|
|
|
} catch { }
|
|
|
|
if (tokenizerConfig) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (tokenizerConfig) {
|
|
if (tokenizerConfig.chat_template) {
|
|
tokenizerConfig.chat_template = formatTemplate(tokenizerConfig.chat_template, tokenizerConfig);
|
|
}
|
|
const info: TokenizerInfo = [tokenizerConfig, tokenizerJson];
|
|
|
|
if (!configOnly) {
|
|
tokenizerCache[modelName] = info;
|
|
saveObject(TOKENIZER_CACHE_KEY, tokenizerCache);
|
|
}
|
|
return info;
|
|
}
|
|
|
|
console.log(`[huggingface] not found config for '${modelName}'`);
|
|
return [null, null];
|
|
};
|
|
|
|
function updateRequired<T extends IParameter>(param: T): T {
|
|
if ('items' in param) {
|
|
updateRequired(param.items);
|
|
} else if ('properties' in param) {
|
|
for (const prop of Object.values(param.properties)) {
|
|
updateRequired(prop);
|
|
}
|
|
param.required = Object.keys(param.properties);
|
|
}
|
|
|
|
return param;
|
|
}
|
|
|
|
const convertFunctionToTool = (fn: IFunction): ITool => ({
|
|
type: 'function',
|
|
function: {
|
|
name: fn.name,
|
|
description: fn.description,
|
|
parameters: updateRequired({
|
|
type: 'object',
|
|
properties: fn.parameters ?? {},
|
|
})
|
|
}
|
|
})
|
|
|
|
export const testToolCalls = (template: string): boolean => {
|
|
const history: ITemplateMessage[] = [
|
|
{ role: 'system', content: 'You are calculator.' },
|
|
{ role: 'user', content: 'Calculate 2 + 2.' },
|
|
];
|
|
|
|
const needle = '___AWOORWA_NEEDLE__';
|
|
|
|
const tools: IFunction[] = [{
|
|
name: 'add',
|
|
description: 'Test function',
|
|
parameters: {
|
|
a: { type: 'number' },
|
|
b: { type: 'number' },
|
|
c: { type: 'array', items: { type: 'number' } },
|
|
d: { type: 'object', properties: { inside: { type: 'number', description: needle } } },
|
|
}
|
|
}];
|
|
|
|
const text = applyChatTemplate(template, history, tools);
|
|
|
|
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 '';
|
|
|
|
let template = templateCache[modelName] ?? null;
|
|
|
|
if (template) {
|
|
console.log(`[huggingface] found cached template for '${modelName}'`);
|
|
} else {
|
|
const [config] = await loadHuggingfaceTokenizer(modelName, true);
|
|
|
|
if (config?.chat_template?.trim()) {
|
|
template = config.chat_template;
|
|
}
|
|
}
|
|
|
|
templateCache[modelName] = template;
|
|
saveObject(TEMPLATE_CACHE_KEY, templateCache);
|
|
|
|
return template;
|
|
}
|
|
|
|
export const findTokenizer = async (modelName: string): Promise<Tokenizer | null> => {
|
|
modelName = normalizeModel(modelName);
|
|
|
|
if (!modelName) return null;
|
|
|
|
let tokenizer = compiledTokenizers.get(modelName) ?? null;
|
|
|
|
if (!tokenizer) {
|
|
const [tokenizerConfig, tokenizerJson] = await loadHuggingfaceTokenizer(modelName);
|
|
if (tokenizerConfig && tokenizerJson) {
|
|
tokenizer = new Tokenizer(tokenizerJson, tokenizerConfig);
|
|
compiledTokenizers.set(modelName, tokenizer);
|
|
}
|
|
}
|
|
|
|
return tokenizer;
|
|
}
|
|
|
|
export const applyChatTemplate = (templateString: string, messages: ITemplateMessage[], functions?: IFunction[]) => (
|
|
applyTemplate(templateString, {
|
|
messages,
|
|
add_generation_prompt: true,
|
|
enable_thinking: false,
|
|
tools: functions?.map(convertFunctionToTool),
|
|
})
|
|
);
|
|
|
|
export const applyTemplate = (templateString: string, args: Record<string, any>): string => {
|
|
try {
|
|
let template = compiledTemplates.get(templateString);
|
|
if (!template) {
|
|
template = new Template(templateString);
|
|
compiledTemplates.set(templateString, template);
|
|
}
|
|
|
|
const result = template.render(args);
|
|
|
|
return result;
|
|
} catch (e) {
|
|
console.error('[applyTemplate] error:', e);
|
|
}
|
|
|
|
return '';
|
|
}
|
|
} |