benchmark.ts

import { gguf, type GGUFTypedMetadata, buildGgufHeader } from "@huggingface/gguf";
import {
  commitIter,
  downloadFile,
  type CommitOutput,
  type CommitProgressEvent,
} from "@huggingface/hub";
import * as fs from "fs";

async function fetchMetadata(
  repoName: string,
  path: string
): Promise<{
  typedMetadata: GGUFTypedMetadata;
  ggufData: {
    tensorDataOffset: number;
    littleEndian: boolean;
    tensorInfoByteRange?: [number, number];
  };
}> {
  const url = `https://huggingface.co/${repoName}/resolve/main/${path}`;
  const ggufResult = await gguf(url, {
    typedMetadata: true,
    additionalFetchHeaders: {
      Authorization: `Bearer ${process.env.HF_TOKEN}`,
    },
  });

  const typedMetadata = ggufResult.typedMetadata;
  const ggufData = {
    tensorDataOffset: Number(ggufResult.tensorDataOffset),
    littleEndian: Boolean(ggufResult.littleEndian ?? true),
    tensorInfoByteRange: ggufResult.tensorInfoByteRange,
  };
  return { typedMetadata, ggufData };
}

async function commitGgufMetadata({
  repoName,
  path,
  metadata,
  ggufData,
}: {
  repoName: string;
  path: string;
  metadata: GGUFTypedMetadata;
  ggufData: {
    tensorDataOffset: number;
    littleEndian: boolean;
    tensorInfoByteRange?: [number, number];
  };
}): Promise<{ commit: CommitOutput["commit"]; pullRequestUrl?: string } | undefined> {
  let uploadProgress = 0;
  let maxUploadProgress = 0;

  let fileUploadProgress: {
    fileName: string;
    fileSizeBytes?: number;
    progressRatio: number;
    state: "uploading" | "error" | "completed";
  } = {
    fileName: path,
    progressRatio: 0,
    state: "uploading", // Show as uploading from the start
  };

  // Reset max progress tracker
  maxUploadProgress = 0;

  try {
    // Reuse existing data
    const tensorDataOffset = ggufData!.tensorDataOffset;
    const littleEndian = ggufData!.littleEndian;
    const tensorInfoByteRange = ggufData!.tensorInfoByteRange;

    if (!tensorInfoByteRange) {
      throw new Error("tensorInfoByteRange not found");
    }

    // Get streaming blob reference - no need to download the full file upfront
    // The commit process will handle downloading only the parts it needs via SplicedBlob
    const originalStreamingBlob = await downloadFile({
      repo: repoName,
      path,
      xet: true, // Enable xet for streaming support
      accessToken: process.env.HF_TOKEN,
    });

    if (!originalStreamingBlob) {
      throw new Error("Failed to get GGUF file reference");
    }

    // Set file size for progress display
    fileUploadProgress.fileSizeBytes = originalStreamingBlob.size;

    const alignment = Number(metadata["general.alignment"]?.value ?? 32);

    const finalHeaderBlob = await buildGgufHeader(originalStreamingBlob, metadata, {
      littleEndian,
      tensorInfoByteRange,
      alignment,
    });

    // Call commitIter with progress tracking
    const commitProgress = commitIter({
      repo: repoName,
      title: "benchmark",
      isPullRequest: true,
      accessToken: process.env.HF_TOKEN, // hf.co/settings/tokens
      // fetch:         fetchWithCredentials,
      useXet: true,
      operations: [
        {
          operation: "edit",
          path,
          originalContent: originalStreamingBlob,
          edits: [
            {
              start: 0,
              end: Number(tensorDataOffset),
              content: finalHeaderBlob,
            },
          ],
        },
      ],
    });

    // Track commit progress
    let result: CommitOutput | undefined;
    let done = false;
    let lastLoggedProgress = -1;
    const progressThreshold = 5; // Log every 5% progress

    while (!done) {
      const { value, done: isDone } = await commitProgress.next();
      done = isDone ?? false;

      if (done) {
        result = value as CommitOutput;
      } else {
        // Handle progress events - value is CommitProgressEvent here
        const progressEvent = value as CommitProgressEvent;
        if (progressEvent.event === "fileProgress") {
          const progressRatio = progressEvent.progress ?? 0;
          const state = progressEvent.state ?? "";
          uploadProgress = progressRatio * 100;

          // Handle the different states with updated progress ranges:
          // - hashing: 0-50% (this is now the download/processing phase)
          // - uploading: 50-100% (actual upload phase)
          if (fileUploadProgress) {
            if (state === "hashing") {
              // During hashing (download/processing), map to 0-50% of total progress
              const newProgressRatio = progressRatio * 0.5;
              if (newProgressRatio > maxUploadProgress) {
                maxUploadProgress = newProgressRatio;
                fileUploadProgress.progressRatio = newProgressRatio;
              }
            } else if (state === "uploading") {
              // During uploading, map to 50-100% of total progress
              const newProgressRatio = 0.5 + progressRatio * 0.5;
              if (newProgressRatio > maxUploadProgress) {
                maxUploadProgress = newProgressRatio;
                fileUploadProgress.progressRatio = newProgressRatio;
              }

              // When upload completes, ensure we're at 100%
              if (progressRatio === 1) {
                fileUploadProgress.progressRatio = 1;
                maxUploadProgress = 1;
              }
            }
          }

          const progress = Math.round(uploadProgress);
          // Only log progress at significant intervals
          if (
            progress >= lastLoggedProgress + progressThreshold ||
            (progress === 100 && lastLoggedProgress < 100)
          ) {
            lastLoggedProgress = progress;
          }
        }
      }
    }

    const hubCommit = result?.commit ?? { oid: "", url: "" };
    const pullRequestUrl = result?.pullRequestUrl;

    // Mark as completed
    if (fileUploadProgress) {
      fileUploadProgress.state = "completed";
      fileUploadProgress.progressRatio = 1;
    }

    // Operation completed successfully

    return { commit: hubCommit, pullRequestUrl };
  } catch (err) {
    // Mark as error
    if (fileUploadProgress && err instanceof Error && err.name !== "AbortError") {
      fileUploadProgress.state = "error";
    }

    throw err;
  }
}

async function getPaths(repoName: string): Promise<string[]> {
  const res = await fetch(`https://huggingface.co/api/models/${repoName}/tree/main`);
  const data = await res.json();
  const paths = data.map((item: any) => item.path);
  const ggufPaths = paths.filter((path: string) => path.endsWith(".gguf"));

  // Sort paths by the numerical value in the filename
  ggufPaths.sort((a: string, b: string) => {
    const numA = parseInt(a.match(/(\d+)mb\.gguf$/)?.[1] || "0");
    const numB = parseInt(b.match(/(\d+)mb\.gguf$/)?.[1] || "0");
    return numA - numB;
  });

  return ggufPaths;
}

(async () => {
  const repoName = "mishig/xet-gguf-edit-test";
  const csvFilePath = "benchmark-results.csv";

  // Initialize CSV file with headers
  fs.writeFileSync(csvFilePath, "path,time_seconds,time_minutes\n");

  const paths = await getPaths(repoName);

  for (const path of paths) {
    console.log(`\nProcessing: ${path}`);

    try {
      const { typedMetadata, ggufData } = await fetchMetadata(repoName, path);
      const newMetadata = {
        ...typedMetadata,
        "tokenizer.chat_template": {
          type: typedMetadata["tokenizer.chat_template"]?.type,
          value: "You are a helpful assistant.",
        },
      } as GGUFTypedMetadata;

      const startTime = performance.now();

      const result = await commitGgufMetadata({
        repoName,
        path,
        metadata: newMetadata,
        ggufData,
      });

      const endTime = performance.now();
      const timeInSeconds = ((endTime - startTime) / 1000).toFixed(2);
      const timeInMinutes = (parseFloat(timeInSeconds) / 60).toFixed(2);

      // Write result to CSV
      fs.appendFileSync(csvFilePath, `${path},${timeInSeconds},${timeInMinutes}\n`);

      console.log(`✓ ${path} - ${timeInSeconds}s (${timeInMinutes}min)`);
      console.log(`  Commit: ${result?.commit.url || "N/A"}`);
    } catch (error) {
      console.error(`✗ ${path} Error:`, error);
    }
  }

  console.log(`\n✓ Benchmark complete! Results saved to ${csvFilePath}`);
})();