Parse tensor infos + E2E test

.eslintrc.json

@@ -5,6 +5,7 @@
  "extends": [
  "eslint:recommended"
  ],
+ "globals": { "fetch": false },
  "ignorePatterns": [
  "coverage/",
  "node_modules/"

README.md

@@ -41,9 +41,9 @@ const arrayBuffer = buffer.buffer.slice(buffer.byteOffset, buffer.byteOffset + b
 const metadata = ggufMetadata(arrayBuffer)
 ```
 
-If you're in a browser environment, you'll probably get parquet file data from either a drag-and-dropped file from the user, or downloaded from the web.
+If you're in a browser environment, you'll probably get .gguf file data from either a drag-and-dropped file from the user, or downloaded from the web.
 
-To load parquet data in the browser from a remote server using `fetch`:
+To load .gguf data in the browser from a remote server using `fetch`:
 
 ```js
 import { ggufMetadata } from 'hyllama'

src/hyllama.d.ts

@@ -4,4 +4,13 @@
  * @param arrayBuffer gguf file contents
  * @returns metadata object
  */
-export declare function ggufMetadata(arrayBuffer: ArrayBuffer): Record<string, any>
+export declare function ggufMetadata(arrayBuffer: ArrayBuffer): {
+ metadata: Record<string, any>;
+ tensorInfos: {
+ name: string;
+ nDims: number;
+ shape: bigint[];
+ type: number;
+ offset: bigint;
+ }[];
+}

src/hyllama.js

@@ -102,7 +102,38 @@ export function ggufMetadata(arrayBuffer) {
  metadata[keyResult.string] = valueResult.value
  }
 
- return metadata
+ const tensorInfos = []
+
+ for (let i = 0; i < tensorCount; i++) {
+ // read tensor name
+ const keyResult = readString(offset)
+ offset = keyResult.newOffset
+
+ const nDims = view.getUint32(offset, true)
+ offset += 4
+
+ /** @type bigint[] */
+ const shape = []
+ for (let dim = 0; dim < nDims; dim++) {
+ shape.push(view.getBigUint64(offset, true))
+ offset += 8
+ }
+
+ const type = view.getUint32(offset, true)
+ offset += 4
+ const tensorOffset = view.getBigUint64(offset, true)
+ offset += 8
+
+ tensorInfos.push({
+ name: keyResult.string,
+ nDims,
+ shape,
+ type,
+ offset: tensorOffset,
+ })
+ }
+
+ return { metadata, tensorInfos }
 }
 
 /**

test/hyllama.test.js

@@ -1,6 +1,8 @@
 import { describe, expect, it } from 'vitest'
 import { ggufMetadata } from '../src/hyllama.js'
 
+const URL_LLAMA = 'https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q2_K.gguf'
+
 describe('ggufMetadata function', () => {
  it('parses metadata correctly', () => {
  // In-memory gguf file
@@ -26,7 +28,7 @@ describe('ggufMetadata function', () => {
  view.setUint32(32 + key.length, 4, true) // value type (UINT32)
  view.setUint32(36 + key.length, value, true) // value
 
- const metadata = ggufMetadata(buffer)
+ const { metadata } = ggufMetadata(buffer)
 
  // Assertions
  expect(metadata).toHaveProperty(key, value)
@@ -40,4 +42,45 @@ describe('ggufMetadata function', () => {
 
  expect(() => ggufMetadata(buffer)).toThrow('Not a valid GGUF file')
  })
+
+ it('parses metadata and tensor info from remote file', async () => {
+ const buf = await (
+ await fetch(URL_LLAMA, {
+ headers: {
+ Range: 'bytes=0-999999',
+ },
+ })
+ ).arrayBuffer()
+
+ const { metadata, tensorInfos } = ggufMetadata(buf)
+
+ expect(metadata).toMatchObject({
+ version: 2,
+ tensorCount: 291,
+ 'general.architecture': 'llama',
+ 'general.file_type': 10,
+ 'general.name': 'LLaMA v2',
+ 'general.quantization_version': 2,
+ 'llama.attention.head_count': 32,
+ 'llama.attention.head_count_kv': 32,
+ 'llama.attention.layer_norm_rms_epsilon': 9.999999974752427e-7,
+ 'llama.block_count': 32,
+ 'llama.context_length': 4096,
+ 'llama.embedding_length': 4096,
+ 'llama.feed_forward_length': 11008,
+ 'llama.rope.dimension_count': 128,
+ })
+
+ expect(tensorInfos.length).toEqual(291)
+ expect(tensorInfos[0]).toMatchObject({
+ name: 'token_embd.weight',
+ shape: [4096n, 32000n],
+ type: 10,
+ })
+ expect(tensorInfos[tensorInfos.length - 1]).toMatchObject({
+ name: 'output_norm.weight',
+ shape: [4096n],
+ type: 0,
+ })
+ })
 })