Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Gguf updates #543

Merged
merged 4 commits into from
Mar 12, 2024
Merged

Gguf updates #543

Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fcab2c9
Use `length` rather than `newOffset`
mishig25 Mar 11, 2024
18f93f3
custom fetch fn
mishig25 Mar 11, 2024
220f41f
format
mishig25 Mar 11, 2024
682483a
params.fetch
mishig25 Mar 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
55 changes: 29 additions & 26 deletions packages/gguf/src/gguf.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,10 @@ class RangeView {

readonly view: DataView;

constructor(public url: string) {
constructor(
public url: string,
private _fetch: typeof fetch = fetch
) {
this.chunk = 0;
/// TODO(fix typing)
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
Expand All @@ -78,7 +81,7 @@ class RangeView {
const range = [this.chunk * HTTP_CHUNK_SIZE, (this.chunk + 1) * HTTP_CHUNK_SIZE - 1];
const buf = new Uint8Array(
await (
await fetch(this.url, {
await this._fetch(this.url, {
headers: {
Range: `bytes=${range[0]}-${range[1]}`,
},
Expand Down Expand Up @@ -115,54 +118,54 @@ function readVersionedSize(view: DataView, byteOffset: number, version: Version)
}
}

function readString(view: DataView, offset: number): { value: string; newOffset: number } {
function readString(view: DataView, offset: number): { value: string; length: number } {
const length = view.getBigUint64(offset, true);
const value = new TextDecoder().decode(view.buffer.slice(offset + 8, offset + 8 + Number(length)));
return { value, newOffset: offset + 8 + Number(length) };
return { value, length: 8 + Number(length) };
}

function readMetadataValue(
view: DataView,
type: GGUFValueType,
offset: number
): { value: MetadataValue; newOffset: number } {
): { value: MetadataValue; length: number } {
switch (type) {
case GGUFValueType.UINT8:
return { value: view.getUint8(offset), newOffset: offset + 1 };
return { value: view.getUint8(offset), length: 1 };
case GGUFValueType.INT8:
return { value: view.getInt8(offset), newOffset: offset + 1 };
return { value: view.getInt8(offset), length: 1 };
case GGUFValueType.UINT16:
return { value: view.getUint16(offset, true), newOffset: offset + 2 };
return { value: view.getUint16(offset, true), length: 2 };
case GGUFValueType.INT16:
return { value: view.getInt16(offset, true), newOffset: offset + 2 };
return { value: view.getInt16(offset, true), length: 2 };
case GGUFValueType.UINT32:
return { value: view.getUint32(offset, true), newOffset: offset + 4 };
return { value: view.getUint32(offset, true), length: 4 };
case GGUFValueType.INT32:
return { value: view.getInt32(offset, true), newOffset: offset + 4 };
return { value: view.getInt32(offset, true), length: 4 };
case GGUFValueType.FLOAT32:
return { value: view.getFloat32(offset, true), newOffset: offset + 4 };
return { value: view.getFloat32(offset, true), length: 4 };
case GGUFValueType.BOOL:
return { value: view.getUint8(offset) !== 0, newOffset: offset + 1 };
return { value: view.getUint8(offset) !== 0, length: 1 };
case GGUFValueType.STRING:
return readString(view, offset);
case GGUFValueType.ARRAY: {
const arrayType = view.getUint32(offset, true);
const arrayLength = view.getBigUint64(offset + 4, true);
let arrayOffset = offset + 12;
let length = 12;
const arrayValues: MetadataValue[] = [];
for (let i = 0; i < arrayLength; i++) {
const { value, newOffset } = readMetadataValue(view, arrayType, arrayOffset);
const { value, length: _length } = readMetadataValue(view, arrayType, offset + length);
arrayValues.push(value);
arrayOffset = newOffset;
length += _length;
}
return { value: arrayValues, newOffset: arrayOffset };
return { value: arrayValues, length };
}
case GGUFValueType.UINT64:
return { value: view.getBigUint64(offset, true), newOffset: offset + 8 };
return { value: view.getBigUint64(offset, true), length: 8 };
case GGUFValueType.INT64:
return { value: view.getBigInt64(offset, true), newOffset: offset + 8 };
return { value: view.getBigInt64(offset, true), length: 8 };
case GGUFValueType.FLOAT64:
return { value: view.getFloat64(offset, true), newOffset: offset + 8 };
return { value: view.getFloat64(offset, true), length: 8 };
}
}

Expand All @@ -185,8 +188,8 @@ export interface GGUFParseOutput {
tensorInfos: GGUFTensorInfo[];
}

export async function gguf(url: string): Promise<GGUFParseOutput> {
const r = new RangeView(url);
export async function gguf(url: string, _fetch: typeof fetch = fetch): Promise<GGUFParseOutput> {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
export async function gguf(url: string, _fetch: typeof fetch = fetch): Promise<GGUFParseOutput> {
export async function gguf(url: string, params: {
/**
* Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
*/
fetch?: typeof fetch;
}): Promise<GGUFParseOutput> {

any way to pass it like this for consistency with safetensors?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

682483a

const r = new RangeView(url, _fetch);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same here

await r.fetchChunk();

if (r.view.getUint32(0, true) !== new DataView(ggufMagicNumber.buffer).getUint32(0, true)) {
Expand All @@ -213,7 +216,7 @@ export async function gguf(url: string): Promise<GGUFParseOutput> {

// read key
const keyResult = readString(r.view, offset);
offset = keyResult.newOffset;
offset += keyResult.length;

// read value type
const valueType = r.view.getUint32(offset, true);
Expand All @@ -223,7 +226,7 @@ export async function gguf(url: string): Promise<GGUFParseOutput> {
throw new Error("Unsupported metadata type: " + valueType);
}

let valueResult: { value: MetadataValue; newOffset: number } | undefined;
let valueResult: ReturnType<typeof readMetadataValue> | undefined;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

while (!valueResult) {
try {
// read value
Expand All @@ -236,7 +239,7 @@ export async function gguf(url: string): Promise<GGUFParseOutput> {
}
}
}
offset = valueResult.newOffset;
offset += valueResult.length;
metadata[keyResult.value] = valueResult.value;
}

Expand All @@ -247,7 +250,7 @@ export async function gguf(url: string): Promise<GGUFParseOutput> {

// read tensor name
const keyResult = readString(r.view, offset);
offset = keyResult.newOffset;
offset += keyResult.length;

const nDims = r.view.getUint32(offset, true);
offset += 4;
Expand Down