206 lines · 5.9 KB · TypeScript
RawDownload
// CRITICALimport { statSync, existsSync, readdirSync, readFileSync } from "node:fs";import { join } from "node:path"; /** * Model metadata information. */export interface ModelInfo {  name: string;  path: string;  size_bytes: number | null;  modified_at: number | null;  architecture: string | null;  quantization: string | null;  context_length: number | null;  recipe_ids: string[];  has_recipe: boolean;} const WEIGHT_EXTS = [".safetensors", ".bin", ".gguf"] as const;const CONFIG_FILENAMES = ["config.json"] as const; /** * Check if a directory looks like a model directory. * @param path - Directory path. * @returns True if it appears to be a model directory. */export const looksLikeModelDirectory = (path: string): boolean => {  if (!existsSync(path)) {    return false;  }  try {    const entries = readdirSync(path, { withFileTypes: true });    for (const configName of CONFIG_FILENAMES) {      if (entries.some((entry) => entry.isFile() && entry.name === configName)) {        return true;      }    }    return entries.some(      (entry) => entry.isFile() && WEIGHT_EXTS.some((extension) => entry.name.toLowerCase().endsWith(extension)),    );  } catch {    return false;  }}; /** * Infer quantization from model name. * @param name - Model directory name. * @returns Quantization identifier. */export const inferQuantization = (name: string): string | undefined => {  const lower = name.toLowerCase();  const candidates = ["awq", "gptq", "gguf", "fp16", "bf16", "int8", "int4", "w4a16", "w8a16"];  return candidates.find((value) => lower.includes(value));}; /** * Read config metadata from config.json. * @param modelDirectory - Model directory. * @returns Metadata object. */export const readConfigMetadata = (modelDirectory: string): { architecture: string | null; context_length: number | null } => {  const configPath = join(modelDirectory, "config.json");  if (!existsSync(configPath)) {    return { architecture: null, context_length: null };  }  try {    const content = readFileSync(configPath, "utf-8");    const parsed = JSON.parse(content) as Record<string, unknown>;    const architectures = parsed["architectures"];    const architecture = Array.isArray(architectures) && architectures.length > 0      ? String(architectures[0])      : null;    const contextLengthRaw =      parsed["max_position_embeddings"] ||      parsed["max_seq_len"] ||      parsed["seq_length"] ||      parsed["n_ctx"];    const contextLength = typeof contextLengthRaw === "number"      ? contextLengthRaw      : typeof contextLengthRaw === "string" && /^\d+$/.test(contextLengthRaw)        ? Number(contextLengthRaw)        : null;    return { architecture, context_length: contextLength };  } catch {    return { architecture: null, context_length: null };  }}; /** * Estimate weight size for a model directory. * @param modelDirectory - Model directory. * @param recursive - Whether to scan recursively. * @returns Total size in bytes. */export const estimateWeightsSizeBytes = (modelDirectory: string, recursive: boolean): number | null => {  let total = 0;  try {    const entries = readdirSync(modelDirectory, { withFileTypes: true });    for (const entry of entries) {      if (entry.isDirectory() && recursive) {        const nested = estimateWeightsSizeBytes(join(modelDirectory, entry.name), true);        total += nested ?? 0;        continue;      }      if (!entry.isFile()) {        continue;      }      if (!WEIGHT_EXTS.some((extension) => entry.name.toLowerCase().endsWith(extension))) {        continue;      }      try {        const stats = statSync(join(modelDirectory, entry.name));        total += stats.size;      } catch {        continue;      }    }  } catch {    return null;  }  return total || null;}; /** * Discover model directories under provided roots. * @param roots - Root paths. * @param maxDepth - Maximum depth. * @param maxModels - Maximum models to return. * @returns List of model directory paths. */export const discoverModelDirectories = (  roots: string[],  maxDepth = 1,  maxModels = 500,): string[] => {  const discovered: string[] = [];  const seen = new Set<string>();  const queue: Array<{ path: string; depth: number }> = roots    .filter((root) => Boolean(root))    .map((root) => ({ path: root, depth: 0 }));   while (queue.length > 0 && discovered.length < maxModels) {    const entry = queue.shift();    if (!entry) {      break;    }    const current = entry.path;    if (seen.has(current)) {      continue;    }    seen.add(current);     if (looksLikeModelDirectory(current)) {      discovered.push(current);      continue;    }     if (entry.depth >= maxDepth) {      continue;    }     try {      const children = readdirSync(current, { withFileTypes: true });      for (const child of children) {        if (!child.isDirectory() || child.name.startsWith(".")) {          continue;        }        queue.push({ path: join(current, child.name), depth: entry.depth + 1 });      }    } catch {      continue;    }  }   return discovered;}; /** * Build model info object for a directory. * @param modelDirectory - Model directory. * @param recipeIds - Associated recipe ids. * @returns Model info. */export const buildModelInfo = async (modelDirectory: string, recipeIds: string[] = []): Promise<ModelInfo> => {  const metadata = await readConfigMetadata(modelDirectory);  let modifiedAt: number | undefined;  try {    modifiedAt = statSync(modelDirectory).mtimeMs;  } catch {    modifiedAt = undefined;  }  const name = modelDirectory.split("/").pop() ?? modelDirectory;  return {    name,    path: modelDirectory,    size_bytes: estimateWeightsSizeBytes(modelDirectory, false),    modified_at: modifiedAt ?? null,    architecture: metadata.architecture,    quantization: inferQuantization(name) ?? null,    context_length: metadata.context_length,    recipe_ids: [...new Set(recipeIds)].sort(),    has_recipe: recipeIds.length > 0,  };};