cast between strict & nonStrict

Author: ngxson

packages/gguf/scripts/generate-llm.ts

@@ -29,7 +29,7 @@ type Attention<TArchitecture extends string> = Record<
 	number
 >>;
 
-type RopeScalingType = "none" | "linear" | "yarn";
+export type TransformerLLMRopeScalingType = "none" | "linear" | "yarn";
 type Rope<TArchitecture extends LLMArchitecture> = Partial<
 	Record<
 			\`\${TArchitecture}.rope.dimension_count\`
@@ -39,7 +39,7 @@ type Rope<TArchitecture extends LLMArchitecture> = Partial<
 		| \`\${TArchitecture}.rope.scaling.original_context_length\`,
 		number
 	>
-	& Record<\`\${TArchitecture}.rope.scaling.type\`, RopeScalingType>
+	& Record<\`\${TArchitecture}.rope.scaling.type\`, TransformerLLMRopeScalingType>
 	& Record<\`\${TArchitecture}.rope.finetuned\`, boolean>
 >;
 

packages/gguf/src/transformer-llm.ts

@@ -19,7 +19,7 @@ type Attention<TArchitecture extends string> = Record<`${TArchitecture}.attentio
 		>
 	>;
 
-type RopeScalingType = "none" | "linear" | "yarn";
+export type TransformerLLMRopeScalingType = "none" | "linear" | "yarn";
 type Rope<TArchitecture extends LLMArchitecture> = Partial<
 	Record<
 		| `${TArchitecture}.rope.dimension_count`
@@ -29,7 +29,7 @@ type Rope<TArchitecture extends LLMArchitecture> = Partial<
 		| `${TArchitecture}.rope.scaling.original_context_length`,
 		number
 	> &
-		Record<`${TArchitecture}.rope.scaling.type`, RopeScalingType> &
+		Record<`${TArchitecture}.rope.scaling.type`, TransformerLLMRopeScalingType> &
 		Record<`${TArchitecture}.rope.finetuned`, boolean>
 >;
 

packages/gguf/src/types.spec.ts

@@ -1,17 +1,30 @@
 import { describe, it } from "vitest";
-import type { GGUFStrictType, GGUFMetadata, GGUFNonStrictType } from "./types";
+import type { gguf } from "./gguf";
+import type { GGUFMetadata, GGUFParseOutput, GGUFType } from "./types";
 
 describe("gguf-types", () => {
-	it("GGUFNonStrictType should be correct (at compile time)", async () => {
+	it("gguf() type can be casted (at compile time)", async () => {
 		// eslint-disable-next-line @typescript-eslint/no-explicit-any
-		const model: GGUFMetadata<GGUFNonStrictType> = null as any;
+		const result: Awaited<ReturnType<typeof gguf>> = null as any;
+		const strictType = result as GGUFParseOutput<GGUFType.strict>;
+		// @ts-expect-error because the key "abc" does not exist
+		strictType.metadata.abc = 123;
+		const nonStrictType = result as GGUFParseOutput<GGUFType.nonStrict>;
+		nonStrictType.metadata.abc = 123; // PASS, because it can be anything
+		// @ts-expect-error because ArrayBuffer is not a MetadataValue
+		nonStrictType.metadata.fff = ArrayBuffer;
+	});
+
+	it("GGUFType.nonStrict should be correct (at compile time)", async () => {
+		// eslint-disable-next-line @typescript-eslint/no-explicit-any
+		const model: GGUFMetadata<GGUFType.nonStrict> = null as any;
 		model.kv_count = 123n;
 		model.abc = 456; // PASS, because it can be anything
 	});
 
-	it("GGUFStrictType should be correct (at compile time)", async () => {
+	it("GGUFType.strict should be correct (at compile time)", async () => {
 		// eslint-disable-next-line @typescript-eslint/no-explicit-any
-		const model: GGUFMetadata<GGUFStrictType> = null as any;
+		const model: GGUFMetadata<GGUFType.strict> = null as any;
 
 		if (model["general.architecture"] === "whisper") {
 			model["encoder.whisper.block_count"] = 0;

packages/gguf/src/types.ts

@@ -110,14 +110,17 @@ export type Whisper = GGUFGeneralInfo<"whisper"> &
 
 /// Types for parse output
 
-export type GGUFStrictType = true;
-export type GGUFNonStrictType = false;
+export enum GGUFType {
+	strict,
+	nonStrict,
+}
 
-export type GGUFMetadata<T extends GGUFStrictType | GGUFNonStrictType = GGUFStrictType> = {
+export type GGUFMetadata<TGGUFType extends GGUFType = GGUFType.strict> = {
 	version: Version;
 	tensor_count: bigint;
 	kv_count: bigint;
-} & (T extends GGUFStrictType ? GGUFModelKV : Record<string, MetadataValue>);
+} & GGUFModelKV &
+	(TGGUFType extends GGUFType.strict ? unknown : Record<string, MetadataValue>);
 
 export type GGUFModelKV = (NoModelMetadata | ModelMetadata) & (NoTokenizer | Tokenizer);
 
@@ -129,7 +132,7 @@ export interface GGUFTensorInfo {
 	offset: bigint;
 }
 
-export interface GGUFParseOutput<T extends GGUFStrictType | GGUFNonStrictType = GGUFStrictType> {
-	metadata: GGUFMetadata<T>;
+export interface GGUFParseOutput<TGGUFType extends GGUFType = GGUFType.strict> {
+	metadata: GGUFMetadata<TGGUFType>;
 	tensorInfos: GGUFTensorInfo[];
 }