diff --git a/src/index.ts b/src/index.ts index 4763aa3..adab5e1 100644 --- a/src/index.ts +++ b/src/index.ts @@ -444,39 +444,19 @@ export const parseRawMetadata = async ( const metadata: Record = {} const setKey = (keyName: string, value: MetadataValue) => { - // this is a bad way to write this and should be clean it up - // but since there are never more than 3 layers currently, it's fine for now - const [key1, key2, key3, key4, key5] = keyName.split('.') - - if (!key2) { - metadata[key1] = value - return - } - if (!key3) { - if (!metadata[key1]) metadata[key1] = {} - metadata[key1][key2] = value - return - } - if (!key4) { - if (!metadata[key1]) metadata[key1] = {} - if (!metadata[key1][key2]) metadata[key1][key2] = {} - metadata[key1][key2][key3] = value - return - } - if (!key5) { - if (!metadata[key1]) metadata[key1] = {} - if (!metadata[key1][key2]) metadata[key1][key2] = {} - if (!metadata[key1][key2][key3]) metadata[key1][key2][key3] = {} - metadata[key1][key2][key3][key4] = value - return - } - if (!metadata[key1]) metadata[key1] = {} - if (!metadata[key1][key2]) metadata[key1][key2] = {} - if (!metadata[key1][key2][key3]) metadata[key1][key2][key3] = {} - if (!metadata[key1][key2][key3][key4]) { - metadata[key1][key2][key3][key4] = {} + const keys = keyName.split('.') + let obj = metadata + for (const [index, key] of keys.entries()) { + if (!index) continue + const prevKey = keys[index - 1] + if (!obj[prevKey]) obj[prevKey] = {} + if (index === keys.length - 1) { + if (typeof obj[prevKey] === 'object') { + obj[prevKey][key] = value + } + } + obj = obj[prevKey] } - metadata[key1][key2][key3][key4][key5] = value } for (let i = 0; i < numKv.value; ++i) { diff --git a/test/index.test.ts b/test/index.test.ts index 7d85924..471b4c6 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -283,4 +283,59 @@ describe('gguf', () => { }, 1000 * 30, ) + + test( + 'DarkForest-20B-v3.0.IQ1_M.gguf', + async () => { + const file = await fetchPartialFile( + 'https://huggingface.co/backyardai/DarkForest-20B-v3.0-GGUF/resolve/main/DarkForest-20B-v3.0.IQ1_M.gguf', + 0, + // 10mb + 1024 * 1024 * 10, + ) + + const fileName = path.join( + __dirname, + 'models', + 'DarkForest-20B-v3.0.IQ1_M.gguf', + ) + + await writeFile(fileName, Buffer.from(file)) + + const { error, metadata } = await gguf(fileName) + + expect(error).toBe(undefined) + expect(metadata).not.toBe(undefined) + if (!metadata) return // for types + + expect(metadata.general.architecture).toBe('llama') + if (!isLlamaMetadata(metadata)) return // for types + + expect(metadata.llama).toBeTruthy() + + expect(metadata).toEqual({ + general: { + architecture: 'llama', + license: 'other', + name: 'F:\\merger\\mergekit\\V3_ultraperecision\\V3_u_4_scaling', + quantization_version: 2, + }, + llama: { + attention: { + head_count: 40, + head_count_kv: 40, + layer_norm_rms_epsilon: 0.000009999999747378752, + }, + context_length: 4096, + embedding_length: 5120, + feed_forward_length: 13824, + rope: { + dimension_count: 128, + freq_base: 10000, + }, + }, + }) + }, + 1000 * 30, + ) })