From b6d0674ce94edb5ed4001e5ddd163dc9d0ffb5b5 Mon Sep 17 00:00:00 2001
From: Douglas Kogut <dkogut1996@gmail.com>
Date: Tue, 30 Jul 2024 17:02:40 -0400
Subject: [PATCH] fix: refactor metadata keying and fix license type

---
 src/index.ts       | 44 ++++++++++---------------------------
 test/index.test.ts | 55 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 32 deletions(-)

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<string, any> = {}
 
   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..8f7ea95 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',
+        'L3-8B-Stheno-v3.3-32K.Q8_0.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,
+  )
 })