@agorapete/wllama 3.5.1-q2.0

package/src/huggingface.ts

@@ -0,0 +1,151 @@
+import { type ModelSource } from './model-manager';
+
+export interface HuggingFaceParams {
+  /**
+   * The repo name, e.g. user/model
+   */
+  repo: string;
+  /**
+   * The file name or path to file in the repo. Only file or quant is needed.
+   */
+  file?: string;
+  /**
+   * The GGUF quantization name, e.g. Q4_K_M, Q8_0, etc. Only file or quant is needed.
+   *
+   * By default, Q4_K_M will be used, then fallback to Q8_0, and finally the non-quantized version if no quantized version is found.
+   */
+  quant?: string;
+  /**
+   * The file name or path to file in the repo for mmproj. Only mmprojFile or mmprojQuant is needed.
+   */
+  mmprojFile?: string;
+  /**
+   * The GGUF quantization name for mmproj, e.g. Q4_K_M, Q8_0, etc. Only mmprojFile or mmprojQuant is needed.
+   */
+  mmprojQuant?: string;
+  /**
+   * The Hugging Face token with permission to access the repo. It can be omitted if the repo is public.
+   */
+  hfToken?: string;
+}
+
+const HF_BASE = 'https://huggingface.co';
+const DEFAULT_QUANTS = ['Q4_K_M', 'Q8_0'];
+
+interface HFFileEntry {
+  type: string;
+  path: string;
+  size: number;
+  oid: string;
+  lfs?: { oid: string; size: number };
+}
+
+async function fetchRepoFiles(
+  repo: string,
+  token?: string
+): Promise<HFFileEntry[]> {
+  const url = `${HF_BASE}/api/models/${repo}/tree/main?recursive=true`;
+  const headers: Record<string, string> = { Accept: 'application/json' };
+  if (token) {
+    headers['Authorization'] = `Bearer ${token}`;
+  }
+  const res = await fetch(url, { headers });
+  if (!res.ok) {
+    let msg = res.statusText;
+    try {
+      msg = (await res.json()).error ?? msg;
+    } catch {
+      /* ignore */
+    }
+    throw new Error(`HF API error (${res.status}): ${msg}`);
+  }
+  return res.json();
+}
+
+// For split GGUF (-00001-of-00005.gguf), return the first shard path.
+// For non-split, return path unchanged.
+function firstShardPath(files: HFFileEntry[], path: string): string {
+  const m = path.match(/^(.+)-(\d{5})-of-(\d{5})\.gguf$/i);
+  if (!m) return path;
+  const first = `${m[1]}-00001-of-${m[3]}.gguf`;
+  return files.some((f) => f.path === first) ? first : path;
+}
+
+function selectFile(
+  files: HFFileEntry[],
+  quant: string | undefined,
+  mmprojOnly: boolean
+): string | null {
+  const candidates = files.filter((f) => {
+    if (f.type !== 'file' || !f.path.toLowerCase().endsWith('.gguf'))
+      return false;
+    const ismmproj = f.path.toLowerCase().includes('mmproj');
+    return mmprojOnly ? ismmproj : !ismmproj;
+  });
+
+  if (candidates.length === 0) return null;
+
+  if (quant) {
+    const upper = quant.toUpperCase();
+    const match = candidates.find((f) => f.path.toUpperCase().includes(upper));
+    if (match) return firstShardPath(candidates, match.path);
+    return null;
+  }
+
+  for (const q of DEFAULT_QUANTS) {
+    const match = candidates.find((f) => f.path.toUpperCase().includes(q));
+    if (match) return firstShardPath(candidates, match.path);
+  }
+
+  // Fallback: first candidate
+  return firstShardPath(candidates, candidates[0].path);
+}
+
+export async function getHFModelSource(
+  config: HuggingFaceParams
+): Promise<ModelSource> {
+  const { repo, file, quant, mmprojFile, mmprojQuant, hfToken } = config;
+
+  const files = await fetchRepoFiles(repo, hfToken);
+
+  const modelPath = file ?? selectFile(files, quant, false);
+  if (!modelPath) {
+    throw new Error(`No GGUF file found in repo "${repo}"`);
+  }
+
+  const source: ModelSource = {
+    url: `${HF_BASE}/${repo}/resolve/main/${modelPath}`,
+  };
+
+  if (mmprojFile || mmprojQuant !== undefined) {
+    const mmpath = mmprojFile ?? selectFile(files, mmprojQuant, true);
+    if (mmpath) {
+      source.mmprojUrl = `${HF_BASE}/${repo}/resolve/main/${mmpath}`;
+    }
+  }
+
+  if (hfToken) {
+    const params = new URLSearchParams({ token: hfToken });
+    source.url += `?${params}`;
+    if (source.mmprojUrl) {
+      source.mmprojUrl += `?${params}`;
+    }
+  }
+
+  return source;
+}
+
+export async function getHFFileSHA256(
+  url: string,
+  headers: Record<string, string>
+): Promise<string | undefined> {
+  if (!url.includes('/resolve/')) return undefined;
+  const rawUrl = url.replace('/resolve/', '/raw/');
+  try {
+    const text = await fetch(rawUrl, { headers }).then((r) => r.text());
+    const match = text.match(/^oid sha256:([0-9a-f]{64})$/m);
+    return match ? match[1] : undefined;
+  } catch {
+    return undefined;
+  }
+}

package/src/index.ts

@@ -0,0 +1,8 @@
+export * from './wllama';
+export * from './cache-manager';
+export * from './model-manager';
+export * from './huggingface';
+export * from './types/types';
+export * from './types/oai-compat';
+export { CacheManager } from './cache-manager';
+export { isValidGgufFile } from './utils';

package/src/mjs.test.ts

@@ -0,0 +1,44 @@
+import { test, expect } from 'vitest';
+import { Wllama as WllamaMJS } from '../esm/index.js';
+import { Wllama as WllamaMJSMinified } from '../esm/index.min.js';
+
+const CONFIG_PATHS = {
+  default: '/src/wasm/wllama.wasm',
+};
+
+const TINY_MODEL =
+  'https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf';
+
+const testFunc = async (wllama: WllamaMJS) => {
+  await wllama.loadModelFromUrl(TINY_MODEL, {
+    n_ctx: 1024,
+  });
+
+  const res = await wllama.createCompletion({
+    prompt: 'Once upon a time',
+    max_tokens: 10,
+    temperature: 0.0,
+    top_p: 0.95,
+    top_k: 40,
+    seed: 42,
+  });
+
+  expect(res).toBeDefined();
+  expect(res.choices[0].text).toMatch(/(there|little|girl|Lily)+/);
+  expect(res.choices[0].text.length).toBeGreaterThan(10);
+
+  await wllama.exit();
+};
+
+// TODO: enable compat mode in tests once test infrastructure supports Safari/asyncify
+test.sequential('(mjs) generates completion', async () => {
+  const wllama = new WllamaMJS(CONFIG_PATHS);
+  wllama.setCompat(null);
+  await testFunc(wllama);
+});
+
+test.sequential('(mjs/minified) generates completion', async () => {
+  const wllama = new WllamaMJSMinified(CONFIG_PATHS);
+  wllama.setCompat(null);
+  await testFunc(wllama as unknown as WllamaMJS);
+});

package/src/model-manager.test.ts

@@ -0,0 +1,200 @@
+import { test, expect } from 'vitest';
+import { ModelManager, Model, ModelValidationStatus } from './model-manager';
+
+const TINY_MODEL =
+  'https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf';
+const SPLIT_MODEL =
+  'https://huggingface.co/ngxson/tinyllama_split_test/resolve/main/stories15M-q8_0-00001-of-00003.gguf';
+
+test.sequential('parseModelUrl handles single model URL', () => {
+  const urls = ModelManager.parseModelUrl(TINY_MODEL);
+  expect(urls.length).toBe(1);
+  expect(urls[0]).toBe(TINY_MODEL);
+});
+
+test.sequential('parseModelUrl handles array of URLs', () => {
+  const urls = ModelManager.parseModelUrl(SPLIT_MODEL);
+  expect(urls.length).toBe(3);
+  expect(urls[0]).toMatch(/-00001-of-00003\.gguf$/);
+  expect(urls[1]).toMatch(/-00002-of-00003\.gguf$/);
+  expect(urls[2]).toMatch(/-00003-of-00003\.gguf$/);
+});
+
+test.sequential('parseModelUrl handles URLs with query parameters', () => {
+  // Test with a simple query parameter
+  const urlWithQuery =
+    'https://example.com/models/model-00001-of-00003.gguf?param=value';
+  const urls = ModelManager.parseModelUrl(urlWithQuery);
+  expect(urls.length).toBe(3);
+  expect(urls[0]).toBe(
+    'https://example.com/models/model-00001-of-00003.gguf?param=value'
+  );
+  expect(urls[1]).toBe(
+    'https://example.com/models/model-00002-of-00003.gguf?param=value'
+  );
+  expect(urls[2]).toBe(
+    'https://example.com/models/model-00003-of-00003.gguf?param=value'
+  );
+
+  // Test with multiple query parameters
+  const urlWithMultipleParams =
+    'https://example.com/models/model-00001-of-00002.gguf?param1=value1&param2=value2';
+  const urlsMultiParams = ModelManager.parseModelUrl(urlWithMultipleParams);
+  expect(urlsMultiParams.length).toBe(2);
+  expect(urlsMultiParams[0]).toBe(
+    'https://example.com/models/model-00001-of-00002.gguf?param1=value1&param2=value2'
+  );
+  expect(urlsMultiParams[1]).toBe(
+    'https://example.com/models/model-00002-of-00002.gguf?param1=value1&param2=value2'
+  );
+
+  // Test with no-inline parameter (common in Vite)
+  const urlWithNoInline =
+    'https://example.com/models/model-00001-of-00002.gguf?no-inline';
+  const urlsNoInline = ModelManager.parseModelUrl(urlWithNoInline);
+  expect(urlsNoInline.length).toBe(2);
+  expect(urlsNoInline[0]).toBe(
+    'https://example.com/models/model-00001-of-00002.gguf?no-inline'
+  );
+  expect(urlsNoInline[1]).toBe(
+    'https://example.com/models/model-00002-of-00002.gguf?no-inline'
+  );
+});
+
+test.sequential('download split model', async () => {
+  const manager = new ModelManager();
+  const model = await manager.downloadModel(SPLIT_MODEL);
+  expect(model.files.length).toBe(3);
+  // check names
+  expect(model.files[0].metadata.originalURL).toMatch(/-00001-of-00003\.gguf$/);
+  expect(model.files[1].metadata.originalURL).toMatch(/-00002-of-00003\.gguf$/);
+  expect(model.files[2].metadata.originalURL).toMatch(/-00003-of-00003\.gguf$/);
+  // check sizes
+  expect(model.files[0].size).toBe(10517152);
+  expect(model.files[1].size).toBe(10381216);
+  expect(model.files[2].size).toBe(5773312);
+});
+
+test.sequential('get downloaded split model', async () => {
+  const manager = new ModelManager();
+  const models = await manager.getModels();
+  const model = models.find((m) => m.url === SPLIT_MODEL);
+  expect(model).toBeDefined();
+  if (!model) throw new Error();
+  // check names
+  expect(model.files[0].metadata.originalURL).toMatch(/-00001-of-00003\.gguf$/);
+  expect(model.files[1].metadata.originalURL).toMatch(/-00002-of-00003\.gguf$/);
+  expect(model.files[2].metadata.originalURL).toMatch(/-00003-of-00003\.gguf$/);
+});
+
+// skip on CI, only run locally with a slow connection
+test.skip('interrupt download split model (partial files downloaded)', async () => {
+  const manager = new ModelManager();
+  await manager.clear();
+  const controller = new AbortController();
+  const downloadPromise = manager.downloadModel(SPLIT_MODEL, {
+    signal: controller.signal,
+    progressCallback: ({ loaded, total }) => {
+      const progress = loaded / total;
+      if (progress > 0.8) {
+        controller.abort();
+      }
+    },
+  });
+  await expect(downloadPromise).rejects.toThrow('aborted');
+  expect((await manager.getModels()).length).toBe(0);
+  expect((await manager.getModels({ includeInvalid: true })).length).toBe(1);
+});
+
+test.sequential('download invalid model URL', async () => {
+  const manager = new ModelManager();
+  const invalidUrl = 'https://invalid.example.com/model.gguf';
+  await expect(manager.downloadModel(invalidUrl)).rejects.toThrow();
+});
+
+test.sequential('download with abort signal', async () => {
+  const manager = new ModelManager();
+  await manager.clear();
+  const controller = new AbortController();
+  const downloadPromise = manager.downloadModel(TINY_MODEL, {
+    signal: controller.signal,
+  });
+  setTimeout(() => controller.abort(), 10);
+  await downloadPromise.catch(console.error);
+  await expect(downloadPromise).rejects.toThrow('aborted');
+  expect((await manager.getModels()).length).toBe(0);
+});
+
+test.sequential('download with progress callback', async () => {
+  const manager = new ModelManager();
+  await manager.clear();
+
+  let progressCalled = false;
+  let lastLoaded = 0;
+  const model = await manager.downloadModel(TINY_MODEL, {
+    progressCallback: ({ loaded, total }) => {
+      expect(loaded).toBeGreaterThan(0);
+      expect(total).toBeGreaterThan(0);
+      expect(loaded).toBeLessThanOrEqual(total);
+      expect(loaded).toBeGreaterThanOrEqual(lastLoaded);
+      progressCalled = true;
+      lastLoaded = loaded;
+    },
+  });
+
+  expect(progressCalled).toBe(true);
+  expect(model).toBeDefined();
+  expect(model.size).toBeGreaterThan(0);
+});
+
+test.sequential('model validation status for new model', async () => {
+  const manager = new ModelManager();
+  const model = new Model(manager, TINY_MODEL);
+  const status = await model.validate();
+  expect(status).toBe(ModelValidationStatus.INVALID);
+});
+
+test.sequential('downloadModel throws on invalid URL', async () => {
+  const manager = new ModelManager();
+  await expect(manager.downloadModel('invalid.txt')).rejects.toThrow();
+});
+
+test.sequential('model size calculation', async () => {
+  const manager = new ModelManager();
+  const model = await manager.downloadModel(TINY_MODEL);
+  expect(model.size).toBe(1185376);
+});
+
+test.sequential('remove model from cache', async () => {
+  const manager = new ModelManager();
+  await manager.clear();
+
+  // Download model first
+  const model = await manager.downloadModel(TINY_MODEL);
+  expect((await manager.getModels()).length).toBe(1);
+  expect(model.size).toBeGreaterThan(0);
+
+  // Remove model
+  await model.remove();
+  expect(model.size).toBe(-1);
+
+  // Try to open removed model
+  await expect(model.open()).rejects.toThrow('deleted from the cache');
+
+  // Validate removed model
+  const status = await model.validate();
+  expect(status).toBe(ModelValidationStatus.DELETED);
+
+  // Cannot see it in list of models
+  const models = await manager.getModels();
+  expect(models.find((m) => m.url === TINY_MODEL)).toBeUndefined();
+});
+
+test.sequential('clear model manager', async () => {
+  const manager = new ModelManager();
+  const model = await manager.downloadModel(TINY_MODEL);
+  expect(model).toBeDefined();
+  expect((await manager.getModels()).length).toBeGreaterThan(0);
+  await manager.clear();
+  expect((await manager.getModels()).length).toBe(0);
+});