gguf: Add types for LLM architectures

packages/gguf/package.json

@@ -32,6 +32,7 @@
 		"format:check": "prettier --check .",
 		"prepublishOnly": "pnpm run build",
 		"build": "tsup src/index.ts --format cjs,esm --clean --dts",
+		"build:llm": "tsx scripts/generate-llm.ts && pnpm run format",
 		"test": "vitest run",
 		"check": "tsc"
 	},

packages/gguf/scripts/generate-llm.ts

@@ -0,0 +1,202 @@
+/**
+ * Script for generating llm.ts
+ * The source data is taken from llama.cpp
+ */
+
+import { writeFileSync } from "node:fs";
+
+const SOURCE_CPP_URL = "https://raw.githubusercontent.com/ggerganov/llama.cpp/master/llama.cpp";
+const DEST_FILE_PATH = "./src/transformer-llm.ts";
+const DEST_COMMON_SOURCE = `
+type Attention<TArchitecture extends string> =
+	& { [K in \`\${TArchitecture}.attention.head_count\`]: number }
+	& { [K in \`\${TArchitecture}.attention.head_count_kv\`]: number }
+	& { [K in \`\${TArchitecture}.attention.layer_norm_epsilon\`]: number }
+	& { [K in \`\${TArchitecture}.attention.layer_norm_rms_epsilon\`]: number }
+	& { [K in \`\${TArchitecture}.attention.alibi_bias_max\`]: number }
+	& { [K in \`\${TArchitecture}.attention.clip_kqv\`]: number }
+	& { [K in \`\${TArchitecture}.attention.use_norm\`]: number };
+
+type Rope<TArchitecture extends LLMArchitecture> =
+	& { [K in \`\${TArchitecture}.rope.dimension_count\`]: number }
+	& { [K in \`\${TArchitecture}.rope.freq_base\`]: number }
+	& { [K in \`\${TArchitecture}.rope.scale\`]: number }
+	& { [K in \`\${TArchitecture}.rope.scale_linear\`]: number };
+
+type MOE<TArchitecture extends LLMArchitecture> =
+	& { [K in \`\${TArchitecture}.expert_count\`]: number }
+	& { [K in \`\${TArchitecture}.expert_used_count\`]: number };
+
+export type TransformerLLMArchitecture = LLMArchitecture; // type alias
+export type TransformerLLMBase<TArchitecture extends LLMArchitecture> = ModelBase<TArchitecture>
+	& MOE<TArchitecture>
+	& Attention<TArchitecture>
+	& Rope<TArchitecture>;
+
+export enum TransformerLLMPoolingType {
+	UNSPECIFIED = -1,
+	NONE = 0,
+	MEAN = 1,
+	CLS = 2,
+};
+`;
+
+const KV_TYPE = {
+	LLM_KV_ATTENTION_LAYERNORM_RMS_EPS: "number",
+	LLM_KV_ATTENTION_LAYERNORM_EPS: "number",
+	LLM_KV_ATTENTION_CAUSAL: "boolean",
+	LLM_KV_TOKENIZER_TOKEN_TYPE_COUNT: "number",
+	LLM_KV_POOLING_TYPE: "TransformerLLMPoolingType",
+	LLM_KV_ATTENTION_CLAMP_KQV: "number",
+	LLM_KV_ATTENTION_MAX_ALIBI_BIAS: "number",
+	LLM_KV_SSM_CONV_KERNEL: "number",
+	LLM_KV_SSM_INNER_SIZE: "number",
+	LLM_KV_SSM_STATE_SIZE: "number",
+	LLM_KV_SSM_TIME_STEP_RANK: "number",
+	LLM_KV_LOGIT_SCALE: "number",
+};
+
+interface Arch {
+	cppConst: string; // for example: "LLM_ARCH_LLAMA"
+	name: string; // for example: "llama"
+	tsName: string; // for example: "ArchLlama"
+	tensorNames: string[]; // for example: "token_embd"
+	hparams: string[];
+}
+
+async function main() {
+	const res = await fetch(SOURCE_CPP_URL);
+	const cppSource = await res.text();
+
+	/////////////////////////////////////
+	// extract list of all architectures
+	const archList: Arch[] = [];
+	const RE_ARCH_NAME = /LLM_ARCH_[A-Z0-9_]+/;
+	const matchedArchList = cppSource.match(/LLM_ARCH_NAMES = (?<names>[^;]+)/)?.groups?.names.split("\n");
+	if (!matchedArchList?.length) {
+		throw new Error("LLM_ARCH_NAMES is empty");
+	}
+	for (const line of matchedArchList) {
+		const matched = line.match(/(?<cppConst>LLM_ARCH_[A-Z0-9_]+),\s+"(?<name>.+?)"/);
+		if (matched?.groups && !matched.groups.name.match(/unknown/)) {
+			archList.push({
+				cppConst: matched.groups.cppConst,
+				name: matched.groups.name,
+				tsName: snakeToPascal(matched.groups.cppConst.replace("LLM_", "")),
+				tensorNames: [],
+				hparams: [],
+			});
+		}
+	}
+
+	/////////////////////////////////////
+	// extract map constant name to kv name
+	// for example: LLM_KV_ATTENTION_LAYERNORM_RMS_EPS ==> "%s.attention.layer_norm_rms_epsilon"
+	const constToKVName: { [cppConst: string]: string } = {};
+	const matchedKVList = cppSource.match(/LLM_KV_NAMES = (?<names>[^;]+)/)?.groups?.names.split("\n");
+	if (!matchedKVList?.length) {
+		throw new Error("LLM_KV_NAMES is empty");
+	}
+	for (const line of matchedKVList) {
+		const matched = line.match(/(?<cppConst>LLM_KV_[A-Z0-9_]+)[,\s]+"(?<name>.+?)"/);
+		if (matched?.groups) {
+			constToKVName[matched.groups.cppConst] = matched.groups.name;
+		}
+	}
+
+	/////////////////////////////////////
+	// extract list of tensor names based on architecture
+	// TODO: unused for now
+	const matchedTensorList = cppSource.match(/LLM_TENSOR_NAMES = (?<names>[^;]+)/)?.groups?.names.split("\n");
+	if (!matchedTensorList?.length) {
+		throw new Error("LLM_TENSOR_NAMES is empty");
+	}
+	let currCppConst = "";
+	for (const line of matchedTensorList) {
+		// check if current line has LLM_ARCH_*
+		const cppConst = line.match(RE_ARCH_NAME)?.[0];
+		if (cppConst) {
+			currCppConst = cppConst;
+			continue;
+		}
+		// check if current line has LLM_TENSOR_*
+		const tensorMatched = line.match(/LLM_TENSOR_[A-Z0-9_]+[,\s]+"(?<name>.+?)"/);
+		if (tensorMatched?.groups) {
+			const arch = archList.find((a) => a.cppConst === currCppConst);
+			if (arch) arch.tensorNames.push(tensorMatched.groups.name);
+		}
+	}
+
+	/////////////////////////////////////
+	// extract list of hyper params based on architecture
+	let insideLoadHParamsFn = false;
+	currCppConst = "";
+	for (const line of cppSource.split("\n")) {
+		// check if current line is function llm_load_hparams()
+		if (line.startsWith("static void llm_load_hparams")) {
+			insideLoadHParamsFn = true;
+		}
+		if (!insideLoadHParamsFn) {
+			continue;
+		}
+		// check if current line has LLM_ARCH_*
+		const RE_CASE = new RegExp(`case (${RE_ARCH_NAME.source})`);
+		const cppConst = line.match(RE_CASE)?.[1];
+		if (cppConst) {
+			currCppConst = cppConst;
+			continue;
+		}
+		// check if current line has get_key(...)
+		const keyConst = line.match(/LLM_KV_[A-Z0-9_]+/)?.[0];
+		if (keyConst) {
+			const arch = archList.find((a) => a.cppConst === currCppConst);
+			if (arch) {
+				arch.hparams.push(keyConst);
+			}
+		}
+		// check if current line is end-of-function
+		if (line === "}") {
+			break;
+		}
+	}
+
+	/////////////////////////////////////
+	// write result to file
+	const content = [
+		"/** This file is auto-generated by generate-llm.ts */",
+		"",
+		'import type { ModelBase } from "./types";',
+		"",
+		"export const LLM_ARCHITECTURES = [",
+		...archList.map((a) => `\t${JSON.stringify(a.name)},`),
+		"] as const;",
+		"type LLMArchitecture = (typeof LLM_ARCHITECTURES)[number];",
+		DEST_COMMON_SOURCE,
+		...archList.map((a) => {
+			let code = `export type ${a.tsName} = TransformerLLMBase<${JSON.stringify(a.name)}>`;
+			if (a.hparams.length) {
+				code += [
+					" & {",
+					...a.hparams.map((k) => `\t${JSON.stringify(constToKVName[k].replace("%s", a.name))}: ${KV_TYPE[k]},`),
+					"};",
+				].join("\n");
+			} else {
+				code += ";";
+			}
+			return code;
+		}),
+		"",
+		`export type TransformerLLM = ${archList.map((a) => a.tsName).join(" | ")};`,
+	].join("\n");
+
+	writeFileSync(DEST_FILE_PATH, content);
+}
+
+function snakeToPascal(str: string) {
+	return str
+		.split("_")
+		.map((word) => word.charAt(0).toUpperCase() + word.slice(1).toLowerCase())
+		.join("");
+}
+
+main();