Merge pull request #5 from ahoylabs/dk/fix-base-schema

fix: make quant version optional and add more options
ahoylabs · Feb 16, 2024 · 43b19c9 · 43b19c9
2 parents e2068af + 22c41b1
commit 43b19c9
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 5 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -368,6 +368,16 @@ const fileTypeIntToString = (
       return 'MOSTLY_Q5_K_M'
     case 18:
       return 'MOSTLY_Q6_K'
+    case 19:
+      return 'MOSTLY_IQ2_XXS'
+    case 20:
+      return 'MOSTLY_IQ2_XS'
+    case 21:
+      return 'MOSTLY_Q2_K_S'
+    case 22:
+      return 'MOSTLY_Q3_K_XS'
+    case 23:
+      return 'MOSTLY_IQ3_XXS'
     default:
       return undefined
   }

diff --git a/src/metadataTypes.ts b/src/metadataTypes.ts
@@ -55,6 +55,11 @@ export type BaseGGUFMetadata = {
     | 'MOSTLY_Q5_K_S'
     | 'MOSTLY_Q5_K_M'
     | 'MOSTLY_Q6_K'
+    | 'MOSTLY_IQ2_XXS'
+    | 'MOSTLY_IQ2_XS'
+    | 'MOSTLY_Q2_K_S'
+    | 'MOSTLY_Q3_K_XS'
+    | 'MOSTLY_IQ3_XXS'
   /**
    * License of the model, expressed as a SPDX license expression
    * (e.g. `"MIT OR Apache-2.0`). *Should not* include any other information,
@@ -75,7 +80,7 @@ export type BaseGGUFMetadata = {
    * scheme's name (e.g. the quantization scheme is Q5_K, and the quantization
    * version is 4).
    **/
-  quantization_version: number
+  quantization_version?: number
   /**
    * Information about where this model came from. This is useful for tracking
    * the provenance of the model, and for finding the original source if the

diff --git a/src/zodValidators.ts b/src/zodValidators.ts
@@ -10,10 +10,9 @@ export const architectureTypeSchema = z.union([
   z.literal('bloom'),
   z.literal('falcon'),
   z.literal('rwkv'),
-  z.string().and(z.object({})),
 ])
 
-const baseGGUFMetadataSchema = z.object({
+export const baseGGUFMetadataSchema = z.object({
   alignment: z.number().optional(),
   author: z.string().optional(),
   description: z.string().optional(),
@@ -38,11 +37,16 @@ const baseGGUFMetadataSchema = z.object({
       z.literal('MOSTLY_Q5_K_S'),
       z.literal('MOSTLY_Q5_K_M'),
       z.literal('MOSTLY_Q6_K'),
+      z.literal('MOSTLY_IQ2_XXS'),
+      z.literal('MOSTLY_IQ2_XS'),
+      z.literal('MOSTLY_Q2_K_S'),
+      z.literal('MOSTLY_Q3_K_XS'),
+      z.literal('MOSTLY_IQ3_XXS'),
     ])
     .optional(),
   license: z.string().optional(),
   name: z.string().optional(),
-  quantization_version: z.number(),
+  quantization_version: z.number().optional(),
   source: z
     .object({
       huggingface: z
@@ -190,7 +194,7 @@ export const falconMetadataSchema = z.object({
     context_length: z.number(),
     embedding_length: z.number(),
     layer_count: z.number(),
-    tensor_data_layout: z.any(),
+    tensor_data_layout: z.string().optional(),
   }),
   general: baseGGUFMetadataSchema.and(
     z.object({

diff --git a/test/index.test.ts b/test/index.test.ts
@@ -178,4 +178,54 @@ describe('gguf', () => {
     },
     1000 * 30,
   )
+
+  test(
+    'EveryoneLLM-7b.gguf',
+    async () => {
+      const file = await fetchPartialFile(
+        'https://huggingface.co/rombodawg/Everyone-LLM-7b-Base-GGUF/resolve/main/EveryoneLLM-7b.gguf',
+        0,
+        // 10mb
+        1024 * 1024 * 10,
+      )
+
+      const fileName = path.join(__dirname, 'models', 'EveryoneLLM-7b.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',
+          file_type: 'MOSTLY_Q8_0',
+          name: 'E:\\Open_source_ai_chatbot\\OOBA_8\\text-generation-webui-main\\models',
+        },
+        llama: {
+          attention: {
+            head_count: 32,
+            head_count_kv: 8,
+            layer_norm_rms_epsilon: 0.000009999999747378752,
+          },
+          context_length: 32768,
+          embedding_length: 4096,
+          feed_forward_length: 14336,
+          rope: {
+            dimension_count: 128,
+            freq_base: 10000,
+          },
+        },
+      })
+    },
+    1000 * 30,
+  )
 })