2024-08-21 14:21:56 +07:00

48 lines
1.4 KiB
TypeScript

import { closeSync, openSync, readSync } from 'fs'
import { Template } from '@huggingface/jinja'
/**
* This is to retrieve the metadata from a GGUF file
* It uses hyllama and jinja from @huggingface module
*/
export const retrieveGGUFMetadata = async (ggufPath: string) => {
try {
const { ggufMetadata } = await import('hyllama')
// Read first 10mb of gguf file
const fd = openSync(ggufPath, 'r')
const buffer = new Uint8Array(10_000_000)
readSync(fd, buffer, 0, 10_000_000, 0)
closeSync(fd)
// Parse metadata and tensor info
const { metadata } = ggufMetadata(buffer.buffer)
const template = new Template(metadata['tokenizer.chat_template'])
const eos_id = metadata['tokenizer.ggml.eos_token_id']
const bos_id = metadata['tokenizer.ggml.bos_token_id']
const eos_token = metadata['tokenizer.ggml.tokens'][eos_id]
const bos_token = metadata['tokenizer.ggml.tokens'][bos_id]
// Parse jinja template
const renderedTemplate = template.render({
add_generation_prompt: true,
eos_token,
bos_token,
messages: [
{
role: 'system',
content: '{system_message}',
},
{
role: 'user',
content: '{prompt}',
},
],
})
return {
...metadata,
parsed_chat_template: renderedTemplate,
}
} catch (e) {
console.log('[MODEL_EXT]', e)
}
}