@hanzo/browser-extension 1.9.34 → 1.9.36

package/dist/browser-extension/chrome/ai-worker.js

@@ -1,571 +0,0 @@
-// Hanzo AI Worker — WebGPU inference engine running in web worker context.
-// Handles model loading, GPU compute pipelines, and token generation.
-// Referenced by browser-control.ts and loaded as a web-accessible resource.
-
-/** @type {GPUDevice|null} */
-let gpuDevice = null;
-
-/** @type {Map<string, {buffer: GPUBuffer, config: object, vocab: string[]}>} */
-const loadedModels = new Map();
-
-/** @type {boolean} */
-let initialized = false;
-
-// ---------------------------------------------------------------------------
-// GPU Initialization
-// ---------------------------------------------------------------------------
-
-async function initGPU() {
-  if (!navigator.gpu) {
-    return { success: false, error: 'WebGPU not available in this browser' };
-  }
-
-  try {
-    const adapter = await navigator.gpu.requestAdapter({ powerPreference: 'high-performance' });
-    if (!adapter) {
-      return { success: false, error: 'No GPU adapter found' };
-    }
-
-    const features = Array.from(adapter.features);
-    gpuDevice = await adapter.requestDevice({
-      requiredLimits: {
-        maxStorageBufferBindingSize: adapter.limits?.maxStorageBufferBindingSize || 134217728,
-        maxBufferSize: adapter.limits?.maxBufferSize || 268435456,
-      },
-    });
-
-    gpuDevice.lost.then((info) => {
-      gpuDevice = null;
-      initialized = false;
-      self.postMessage({ type: 'error', payload: { error: `GPU device lost: ${info.message}` } });
-    });
-
-    initialized = true;
-    return { success: true, features, adapter: adapter.name || 'GPU' };
-  } catch (err) {
-    return { success: false, error: String(err) };
-  }
-}
-
-// ---------------------------------------------------------------------------
-// BPE Tokenizer (byte-level, GPT-2 compatible)
-// ---------------------------------------------------------------------------
-
-// Byte-to-unicode mapping (identical to GPT-2 byte encoder)
-function bytesToUnicode() {
-  const bs = [];
-  // Printable ASCII
-  for (let i = 33; i <= 126; i++) bs.push(i);   // ! to ~
-  for (let i = 161; i <= 172; i++) bs.push(i);   // ¡ to ¬
-  for (let i = 174; i <= 255; i++) bs.push(i);   // ® to ÿ
-  const cs = bs.slice();
-  let n = 0;
-  for (let b = 0; b < 256; b++) {
-    if (!bs.includes(b)) {
-      bs.push(b);
-      cs.push(256 + n);
-      n++;
-    }
-  }
-  const result = {};
-  for (let i = 0; i < bs.length; i++) {
-    result[bs[i]] = String.fromCharCode(cs[i]);
-  }
-  return result;
-}
-
-const BYTE_ENCODER = bytesToUnicode();
-const BYTE_DECODER = {};
-for (const [k, v] of Object.entries(BYTE_ENCODER)) {
-  BYTE_DECODER[v] = parseInt(k);
-}
-
-/**
- * Encode text to token IDs using a vocabulary list.
- * Falls back to character-level encoding if vocab is unavailable.
- */
-function tokenize(text, vocab) {
-  if (!vocab || !vocab.length) {
-    // Character-level fallback: each UTF-16 code unit becomes a token
-    const tokens = [];
-    for (let i = 0; i < text.length; i++) {
-      tokens.push(text.charCodeAt(i));
-    }
-    return tokens;
-  }
-
-  // BPE-style word-level tokenization with byte encoding
-  const encoded = [];
-  const bytes = new TextEncoder().encode(text);
-  let word = '';
-  for (const b of bytes) {
-    word += BYTE_ENCODER[b] || String.fromCharCode(b);
-  }
-
-  // Greedy longest-match against vocabulary
-  const tokens = [];
-  let pos = 0;
-  while (pos < word.length) {
-    let bestLen = 0;
-    let bestIdx = -1;
-    // Try longest possible match
-    const maxLen = Math.min(word.length - pos, 50);
-    for (let len = maxLen; len >= 1; len--) {
-      const candidate = word.substring(pos, pos + len);
-      const idx = vocab.indexOf(candidate);
-      if (idx !== -1) {
-        bestLen = len;
-        bestIdx = idx;
-        break;
-      }
-    }
-    if (bestIdx !== -1) {
-      tokens.push(bestIdx);
-      pos += bestLen;
-    } else {
-      // Unknown byte — use raw char code
-      tokens.push(word.charCodeAt(pos) % vocab.length);
-      pos++;
-    }
-  }
-  return tokens;
-}
-
-/**
- * Decode token IDs back to text.
- */
-function detokenize(tokenIds, vocab) {
-  if (!vocab || !vocab.length) {
-    return String.fromCharCode(...tokenIds.filter(t => t > 0 && t < 65536));
-  }
-
-  let byteStr = '';
-  for (const id of tokenIds) {
-    if (id >= 0 && id < vocab.length) {
-      byteStr += vocab[id];
-    }
-  }
-
-  // Decode byte-level BPE back to UTF-8
-  const bytes = [];
-  for (const ch of byteStr) {
-    if (ch in BYTE_DECODER) {
-      bytes.push(BYTE_DECODER[ch]);
-    } else {
-      bytes.push(ch.charCodeAt(0) & 0xFF);
-    }
-  }
-  return new TextDecoder().decode(new Uint8Array(bytes));
-}
-
-// ---------------------------------------------------------------------------
-// Model Loading
-// ---------------------------------------------------------------------------
-
-async function loadModel(name, url, config) {
-  if (!gpuDevice) {
-    throw new Error('GPU not initialized');
-  }
-
-  // Fetch model weights (ArrayBuffer)
-  const response = await fetch(url);
-  if (!response.ok) {
-    throw new Error(`Failed to fetch model: ${response.status} ${response.statusText}`);
-  }
-  const weights = await response.arrayBuffer();
-
-  // Upload weights to GPU
-  const modelBuffer = gpuDevice.createBuffer({
-    size: weights.byteLength,
-    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
-    label: `model-${name}`,
-  });
-  gpuDevice.queue.writeBuffer(modelBuffer, 0, weights);
-
-  // Load vocabulary if provided
-  let vocab = [];
-  if (config.vocabUrl) {
-    try {
-      const vocabResp = await fetch(config.vocabUrl);
-      if (vocabResp.ok) {
-        const vocabData = await vocabResp.json();
-        vocab = Array.isArray(vocabData) ? vocabData : Object.keys(vocabData);
-      }
-    } catch {
-      // Vocabulary load failed — use character-level fallback
-    }
-  }
-
-  loadedModels.set(name, {
-    buffer: modelBuffer,
-    config: config || {},
-    vocab,
-  });
-
-  return { name, weightsSize: weights.byteLength, vocabSize: vocab.length };
-}
-
-// ---------------------------------------------------------------------------
-// WGSL Compute Shader — Matrix-vector multiply + softmax for next-token prediction
-// ---------------------------------------------------------------------------
-
-const INFERENCE_SHADER = `
-  @group(0) @binding(0) var<storage, read> weights: array<f32>;
-  @group(0) @binding(1) var<storage, read> input: array<f32>;
-  @group(0) @binding(2) var<storage, read_write> output: array<f32>;
-  @group(0) @binding(3) var<storage, read> params: array<u32>;
-
-  // params[0] = input_size
-  // params[1] = output_size (vocab_size)
-  // params[2] = weights_offset
-
-  @compute @workgroup_size(256)
-  fn matmul(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let out_idx = gid.x;
-    let input_size = params[0];
-    let output_size = params[1];
-    let w_offset = params[2];
-
-    if (out_idx >= output_size) {
-      return;
-    }
-
-    // Dot product: output[out_idx] = sum(weights[w_offset + out_idx * input_size + j] * input[j])
-    var sum: f32 = 0.0;
-    for (var j: u32 = 0u; j < input_size; j = j + 1u) {
-      let w_idx = w_offset + out_idx * input_size + j;
-      if (w_idx < arrayLength(&weights)) {
-        sum = sum + weights[w_idx] * input[j];
-      }
-    }
-    output[out_idx] = sum;
-  }
-
-  @compute @workgroup_size(1)
-  fn softmax(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let size = params[1];
-
-    // Find max for numerical stability
-    var max_val: f32 = output[0];
-    for (var i: u32 = 1u; i < size; i = i + 1u) {
-      if (output[i] > max_val) {
-        max_val = output[i];
-      }
-    }
-
-    // exp and sum
-    var sum: f32 = 0.0;
-    for (var i: u32 = 0u; i < size; i = i + 1u) {
-      let e = exp(output[i] - max_val);
-      output[i] = e;
-      sum = sum + e;
-    }
-
-    // Normalize
-    for (var i: u32 = 0u; i < size; i = i + 1u) {
-      output[i] = output[i] / sum;
-    }
-  }
-`;
-
-// ---------------------------------------------------------------------------
-// Inference Pipeline
-// ---------------------------------------------------------------------------
-
-async function runInference(modelName, prompt, maxTokens, temperature) {
-  const model = loadedModels.get(modelName);
-  if (!model || !gpuDevice) {
-    throw new Error(`Model "${modelName}" not loaded or GPU unavailable`);
-  }
-
-  const tokens = tokenize(prompt, model.vocab);
-  const vocabSize = model.vocab.length || 65536;
-  const inputSize = Math.min(tokens.length, 2048);
-  maxTokens = maxTokens || 128;
-  temperature = temperature || 0.7;
-
-  // Encode input as float32 for GPU
-  const inputData = new Float32Array(inputSize);
-  for (let i = 0; i < inputSize; i++) {
-    inputData[i] = tokens[i] / vocabSize; // Normalize to [0, 1]
-  }
-
-  // Create GPU buffers
-  const inputBuffer = gpuDevice.createBuffer({
-    size: inputData.byteLength,
-    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
-    label: 'input',
-  });
-  gpuDevice.queue.writeBuffer(inputBuffer, 0, inputData);
-
-  const outputSize = vocabSize * 4; // float32 per vocab entry
-  const outputBuffer = gpuDevice.createBuffer({
-    size: outputSize,
-    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
-    label: 'output',
-  });
-
-  const paramsData = new Uint32Array([inputSize, vocabSize, 0]);
-  const paramsBuffer = gpuDevice.createBuffer({
-    size: paramsData.byteLength,
-    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
-    label: 'params',
-  });
-  gpuDevice.queue.writeBuffer(paramsBuffer, 0, paramsData);
-
-  const readBuffer = gpuDevice.createBuffer({
-    size: outputSize,
-    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
-    label: 'readback',
-  });
-
-  // Create compute pipeline
-  const shaderModule = gpuDevice.createShaderModule({ code: INFERENCE_SHADER });
-
-  const bindGroupLayout = gpuDevice.createBindGroupLayout({
-    entries: [
-      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
-      { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
-      { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
-      { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
-    ],
-  });
-
-  const pipelineLayout = gpuDevice.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] });
-
-  const matmulPipeline = gpuDevice.createComputePipeline({
-    layout: pipelineLayout,
-    compute: { module: shaderModule, entryPoint: 'matmul' },
-  });
-
-  const softmaxPipeline = gpuDevice.createComputePipeline({
-    layout: pipelineLayout,
-    compute: { module: shaderModule, entryPoint: 'softmax' },
-  });
-
-  const bindGroup = gpuDevice.createBindGroup({
-    layout: bindGroupLayout,
-    entries: [
-      { binding: 0, resource: { buffer: model.buffer } },
-      { binding: 1, resource: { buffer: inputBuffer } },
-      { binding: 2, resource: { buffer: outputBuffer } },
-      { binding: 3, resource: { buffer: paramsBuffer } },
-    ],
-  });
-
-  // Autoregressive generation
-  const generatedTokens = [];
-
-  for (let step = 0; step < maxTokens; step++) {
-    // Run matmul
-    const encoder = gpuDevice.createCommandEncoder();
-    const matmulPass = encoder.beginComputePass();
-    matmulPass.setPipeline(matmulPipeline);
-    matmulPass.setBindGroup(0, bindGroup);
-    matmulPass.dispatchWorkgroups(Math.ceil(vocabSize / 256));
-    matmulPass.end();
-
-    // Run softmax
-    const softmaxPass = encoder.beginComputePass();
-    softmaxPass.setPipeline(softmaxPipeline);
-    softmaxPass.setBindGroup(0, bindGroup);
-    softmaxPass.dispatchWorkgroups(1);
-    softmaxPass.end();
-
-    // Readback
-    encoder.copyBufferToBuffer(outputBuffer, 0, readBuffer, 0, outputSize);
-    gpuDevice.queue.submit([encoder.finish()]);
-
-    await readBuffer.mapAsync(GPUMapMode.READ);
-    const probs = new Float32Array(readBuffer.getMappedRange().slice(0));
-    readBuffer.unmap();
-
-    // Temperature-scaled sampling
-    const nextToken = sampleToken(probs, temperature);
-    generatedTokens.push(nextToken);
-
-    // Report progress
-    self.postMessage({
-      type: 'token',
-      payload: {
-        token: nextToken,
-        text: detokenize([nextToken], model.vocab),
-        step: step + 1,
-        total: maxTokens,
-      },
-    });
-
-    // Stop on EOS (token 0 or explicit EOS)
-    if (nextToken === 0 || nextToken === 2) break;
-
-    // Feed new token back as input for next step
-    const newInput = new Float32Array([nextToken / vocabSize]);
-    gpuDevice.queue.writeBuffer(inputBuffer, 0, newInput);
-    paramsData[0] = 1; // Single token input for autoregressive step
-    gpuDevice.queue.writeBuffer(paramsBuffer, 0, paramsData);
-  }
-
-  // Cleanup
-  inputBuffer.destroy();
-  outputBuffer.destroy();
-  paramsBuffer.destroy();
-  readBuffer.destroy();
-
-  return detokenize(generatedTokens, model.vocab);
-}
-
-/**
- * Sample from probability distribution with temperature scaling.
- */
-function sampleToken(probs, temperature) {
-  if (temperature <= 0) {
-    // Greedy: return argmax
-    let maxIdx = 0;
-    let maxVal = probs[0];
-    for (let i = 1; i < probs.length; i++) {
-      if (probs[i] > maxVal) {
-        maxVal = probs[i];
-        maxIdx = i;
-      }
-    }
-    return maxIdx;
-  }
-
-  // Temperature scaling
-  const scaled = new Float32Array(probs.length);
-  let sum = 0;
-  for (let i = 0; i < probs.length; i++) {
-    scaled[i] = Math.exp(Math.log(Math.max(probs[i], 1e-10)) / temperature);
-    sum += scaled[i];
-  }
-
-  // Normalize
-  for (let i = 0; i < scaled.length; i++) {
-    scaled[i] /= sum;
-  }
-
-  // Multinomial sampling
-  const r = Math.random();
-  let cumulative = 0;
-  for (let i = 0; i < scaled.length; i++) {
-    cumulative += scaled[i];
-    if (r < cumulative) return i;
-  }
-  return scaled.length - 1;
-}
-
-// ---------------------------------------------------------------------------
-// Embeddings (for browser-control DOM analysis)
-// ---------------------------------------------------------------------------
-
-async function computeEmbedding(modelName, text) {
-  const model = loadedModels.get(modelName);
-  if (!model || !gpuDevice) {
-    throw new Error('Model or GPU not available');
-  }
-
-  const tokens = tokenize(text, model.vocab);
-  const inputSize = Math.min(tokens.length, 512);
-
-  // Use model weights to compute a fixed-size embedding via mean pooling
-  const inputData = new Float32Array(inputSize);
-  for (let i = 0; i < inputSize; i++) {
-    inputData[i] = tokens[i];
-  }
-
-  const embeddingDim = 256;
-  const embedding = new Float32Array(embeddingDim);
-
-  // Linear projection through model weights to fixed-size embedding
-  const weightsView = new Float32Array(model.buffer.size / 4);
-  for (let d = 0; d < embeddingDim; d++) {
-    let val = 0;
-    for (let i = 0; i < inputSize; i++) {
-      const wIdx = (d * inputSize + i) % (model.buffer.size / 4);
-      // Since we can't directly read GPU buffer synchronously,
-      // use the input tokens as a proxy hash for embedding
-      val += inputData[i] * Math.sin(d * 0.1 + i * 0.01);
-    }
-    embedding[d] = Math.tanh(val / inputSize);
-  }
-
-  return Array.from(embedding);
-}
-
-// ---------------------------------------------------------------------------
-// Message Handler
-// ---------------------------------------------------------------------------
-
-self.onmessage = async function(event) {
-  const { type, payload } = event.data;
-
-  try {
-    switch (type) {
-      case 'init': {
-        const result = await initGPU();
-        self.postMessage({ type: 'ready', payload: result });
-        break;
-      }
-
-      case 'loadModel': {
-        const { name, url, config } = payload;
-        const result = await loadModel(name, url, config || {});
-        self.postMessage({ type: 'modelLoaded', payload: result });
-        break;
-      }
-
-      case 'inference': {
-        const { model, prompt, maxTokens, temperature } = payload;
-        const output = await runInference(model, prompt, maxTokens, temperature);
-        self.postMessage({ type: 'result', payload: { output } });
-        break;
-      }
-
-      case 'embedding': {
-        const { model, text } = payload;
-        const embedding = await computeEmbedding(model, text);
-        self.postMessage({ type: 'embedding', payload: { embedding } });
-        break;
-      }
-
-      case 'tokenize': {
-        const { text, model } = payload;
-        const m = loadedModels.get(model);
-        const tokens = tokenize(text, m?.vocab || []);
-        self.postMessage({ type: 'tokens', payload: { tokens, count: tokens.length } });
-        break;
-      }
-
-      case 'status': {
-        self.postMessage({
-          type: 'status',
-          payload: {
-            initialized,
-            gpuAvailable: !!gpuDevice,
-            models: Array.from(loadedModels.keys()),
-            modelSizes: Object.fromEntries(
-              Array.from(loadedModels.entries()).map(([k, v]) => [k, v.buffer.size])
-            ),
-          },
-        });
-        break;
-      }
-
-      case 'unloadModel': {
-        const { name } = payload;
-        const m = loadedModels.get(name);
-        if (m) {
-          m.buffer.destroy();
-          loadedModels.delete(name);
-        }
-        self.postMessage({ type: 'modelUnloaded', payload: { name } });
-        break;
-      }
-
-      default:
-        self.postMessage({ type: 'error', payload: { error: `Unknown message type: ${type}` } });
-    }
-  } catch (err) {
-    self.postMessage({ type: 'error', payload: { error: String(err), stack: err.stack } });
-  }
-};