diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 83cd0a1..f95a722 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,6 +79,9 @@ jobs: - name: Run tests with coverage run: npm run test:coverage + - name: Coverage threshold check + run: node scripts/ci-coverage-check.js --tier integration --verbose + - name: Upload coverage artifact if: always() uses: actions/upload-artifact@v7 @@ -178,22 +181,18 @@ jobs: - name: Generate API documentation run: npx typedoc - continue-on-error: true - name: Download benchmark report uses: actions/download-artifact@v4 with: { name: benchmark-report, path: docs/benchmark } - continue-on-error: true - name: Download web benchmark report uses: actions/download-artifact@v4 with: { name: web-benchmark-report, path: docs/web-benchmark } - continue-on-error: true - name: Download coverage report uses: actions/download-artifact@v4 with: { name: coverage-report, path: docs/coverage } - continue-on-error: true - name: Generate combined benchmark report run: | @@ -206,7 +205,6 @@ jobs: else echo "[deploy-pages] Skipping combined benchmark (missing benchmark data)" fi - continue-on-error: true - name: Prepare pages run: node scripts/prepare-pages.js diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 95bd791..4448763 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -42,7 +42,6 @@ jobs: - name: Unit Tests (no ONNX model needed) run: npm run test:unit:coverage - continue-on-error: true - name: Upload unit coverage if: always() @@ -81,12 +80,10 @@ jobs: run: npm run test:coverage env: EMBED_CODE_MODEL_PATH: models/nomic-embed-code-v1.5.int8.onnx - continue-on-error: true - name: Check coverage thresholds if: ${{ !inputs.skip_tests }} run: node scripts/ci-coverage-check.js --tier release --verbose - continue-on-error: true - name: Upload integration coverage if: always() @@ -121,11 +118,10 @@ jobs: models/nomic-embed-code-v1.5.int8.onnx key: embed-code-model-${{ hashFiles("scripts/export-model.py", "models/model-descriptor.json") }}-v5 fail-on-cache-miss: true - fail-on-cache-miss: true - name: Run benchmark run: | - npx tsx scripts/benchmark-ci.js + npx tsx scripts/benchmark-ci.js --verbose --json benchmark-report.json --md benchmark-report.md --html benchmark-report.html timeout-minutes: 10 env: EMBED_CODE_MODEL_PATH: models/nomic-embed-code-v1.5.int8.onnx @@ -133,13 +129,15 @@ jobs: - name: Accuracy gate run: node scripts/ci-benchmark-check.js --verbose - continue-on-error: true - name: Upload benchmark artifact uses: actions/upload-artifact@v7 with: name: benchmark-report - path: benchmarks/ + path: | + benchmark-report.json + benchmark-report.md + benchmark-report.html retention-days: 7 # ═══════════════════════════════════════════════════════════════════ @@ -193,7 +191,6 @@ jobs: contents: read pages: write id-token: write - continue-on-error: true environment: name: github-pages url: ${{ steps.deployment.outputs.page_url }} @@ -209,21 +206,18 @@ jobs: - name: Generate API documentation run: npx typedoc - continue-on-error: true - name: Download benchmark report uses: actions/download-artifact@v4 with: name: benchmark-report path: docs/benchmark - continue-on-error: true - name: Download coverage report uses: actions/download-artifact@v4 with: name: coverage-report path: docs/coverage - continue-on-error: true - name: Prepare pages run: node scripts/prepare-pages.js diff --git a/.prettierignore b/.prettierignore index 877dc30..3ad5eaf 100644 --- a/.prettierignore +++ b/.prettierignore @@ -1,6 +1,7 @@ dist/ coverage/ node_modules/ +docs/ *.tsbuildinfo models/*.onnx models/*.zip diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md deleted file mode 100644 index cbb985c..0000000 --- a/docs/ARCHITECTURE.md +++ /dev/null @@ -1,211 +0,0 @@ -# Architecture - -> `embed-code-ts` — A TypeScript/Node.js implementation of Nomic's nomic-embed-code model with int8 incbin-style model delivery. - ---- - -## High-Level Architecture - -``` -┌──────────────────────────────────────────────────────────────┐ -│ @agentix-e/embed-code-cli │ -│ Commander-based CLI: setup (model download) + embed │ -└──────────────────────────┬───────────────────────────────────┘ - │ uses -┌──────────────────────────▼───────────────────────────────────┐ -│ @agentix-e/embed-code-core │ -│ ┌──────────┐ ┌──────────┐ ┌───────────┐ ┌────────────┐ │ -│ │ Model │ │ Config │ │ Tokenizer │ │ Pooler │ │ -│ │ (API) │ │ (types) │ │ (BPE) │ │ (strategies)│ │ -│ └────┬─────┘ └──────────┘ └─────┬─────┘ └─────┬──────┘ │ -│ │ │ │ │ -│ │ ┌───────────────────────┴───────────────┘ │ -│ │ │ │ -│ ▼ ▼ │ -│ ┌──────────────────────────────────────────┐ │ -│ │ Embedding Pipeline │ │ -│ │ Step 1: Tokenize (BPE) │ │ -│ │ Step 2: -│ │ Step 3: Pool (last_token/mean/cls) │ │ -│ │ Step 4: Normalize (L2) │ │ -│ └────────────────┬─────────────────────────┘ │ -│ │ uses │ -│ ┌────────────────▼─────────────────────────┐ │ -│ │ EmbedCodeInferenceEngine │ │ -│ │ IInferenceEngine → pure-TypeScript backend │ │ -│ │ CPU / CUDA / DirectML │ │ -│ └──────────────────────────────────────────┘ │ -│ │ -│ Utilities: Model Downloader, Descriptor Resolver │ -│ Descriptor: model-descriptor.json (architecture contract) │ -└──────────────────────────────────────────────────────────────┘ -``` - ---- - -## Component Design - -### 1. EmbedCode (Public API) - -**File**: `packages/embed-code-core/src/embed-code.ts` - -The main entry point. Implements the public API via a private constructor + static factory. - -```typescript -// Lifecycle -embedder = EmbedCode.fromPretrained({ modelPath, executionProvider }); -results = embedder.embed(texts, { maxTokens, poolingStrategy, normalize }); -sim = embedder.similarity(embA, embB); -embedder.dispose(); -``` - -**Design decisions**: - -- **Private constructor + static factory**: Ensures async initialization is enforced -- **Task prefix API**: `embedder.taskPrefixes` exposes `{ query, document }` for proper prompt formatting -- **Flat Float32Array output**: Optimized for zero-copy bulk similarity computation - -### 2. Tokenizer Pipeline - -**File**: `packages/embed-code-core/src/tokenizer.ts` - -Transforms raw text into model-ready token IDs: - -``` -Raw Text → [Normalize] → [BPE Encode] → [Pad/Truncate] → [Attention Mask] → [Token IDs] - │ │ │ │ - Unicode NFC Vocabulary maxTokens 1=real, 0=pad - + lowercase lookup -``` - -### 3. Embedding Pipeline - -``` -Input Text(s) - │ - ├─→ [tokenize()] - │ ├─ Text normalization (Unicode NFC) - │ ├─ BPE encode → token IDs - │ ├─ Truncate/pad → maxTokens - │ └─ Generate attention mask - │ - ├─→ [infer()] - │ ├─ Build build input tensors - │ ├─ engine.forward() → hidden states - │ └─ Extract last_hidden_state [B, L, 3584] - │ - └─→ [pool + normalize()] - ├─ Pool (last_token / mean / cls) - ├─ L2 normalize → unit vectors - └─ Return Float32Array(N × 3584) -``` - -### 4. TS Inference Engine - -**File**: `packages/embed-code-core/src/inference/ts-engine.ts` - -Implements `IInferenceEngine`: - -- **Pluggable execution providers**: CPU / CUDA / DirectML -- **Dynamic shape handling**: Variable-length inputs auto-padded to model maxTokens -- **Concurrent batch inference**: `Promise.all` for parallel TS session calls -- **Proper resource cleanup**: Calls `session.release()` on dispose - -### 5. Pooling Strategies - -**File**: `packages/embed-code-core/src/pooling.ts` - -| Strategy | Description | Use Case | -| -------------- | ----------------------------------------------- | ------------------------------- | -| **last_token** | Use the last non-padding token embedding | Default, recommended for search | -| **Mean** | Average of all token embeddings (ignoring pads) | Best for code classification | -| **CLS** | Use the [CLS] token embedding only | BERT-style tasks | - -### 6. Model Downloader - -**File**: `packages/embed-code-core/src/model-downloader.ts` - -- **Streaming download**: Uses Node.js `fetch` reader → file `writeStream` (no large heap buffer for multi-GB models) -- **Proxy support**: Environment variables (`EMBED_CODE_PROXY_URL/USERNAME/PASSWORD`, `HTTPS_PROXY`) or programmatic `DownloadOptions.proxy` with username/password. Uses undici `ProxyAgent` for clean proxy handling without global environment mutations -- **SHA-256 integrity**: Hashes the TS file after download for verification -- **Cache management**: Platform-aware cache directory (`XDG_CACHE_HOME`) -- **Error hierarchy**: `ProxyAuthError` (HTTP 407), `DownloadError`, `ChecksumMismatchError` - -### 7. Model Descriptor System - -**File**: `packages/embed-code-core/src/model-descriptor.ts` - -The `model-descriptor.json` file (committed to repo, distributed with releases) defines the architecture contract: - -``` -model-descriptor.json → resolveModelConfig() → ModelConfig - │ - └─ Single source of truth for: - ├─ model.hf_revision (HuggingFace checkpoint) - ├─ weights.sha256 (download integrity) - ├─ architecture.* (dim, maxTokens, vocabSize) - └─ processing.* (tokenizer config, pooling defaults) -``` - ---- - -## Data Flow - -``` -User Input (string | string[]) - │ - ▼ -[EmbedCode.embed()] - │ - ├─→ [tokenize()] - │ ├─ Normalize: Unicode NFC + lowercase - │ ├─ BPE encode → token IDs - │ ├─ Pad/truncate → maxTokens (32768) - │ └─ Attention mask generation - │ - ├─→ [inference()] - │ ├─ Build BigInt64Array input_ids [B, L] - │ ├─ Build BigInt64Array attention_mask [B, L] - │ ├─ engine.forward({ input_ids, attention_mask }) - │ └─ Extract last_hidden_state [B, L, 3584] - │ - └─→ [pool + normalize()] - ├─ Pool: select strategy (last_token/mean/cls) - ├─ L2 normalize: x / ||x||₂ - └─ Return { embeddings: Float32Array(N×3584), elapsedMs } -``` - ---- - -## Type System - -``` -EmbedOptions — mutable options (passed to embed()) -ModelConfig — frozen, read-only (resolved from descriptor) -EmbeddingResult — { embeddings: Float32Array, shape: [number, number], elapsedMs: number } -EmbedProgress — { phase: 'tokenize' | 'inference' | 'pool' | 'normalize', step: number, total: number } -IInferenceEngine — pluggable backend (TS) -``` - ---- - -## Key Design Principles - -1. **Incbin-inspired delivery**: npm package is code-only (~50 KB); model is downloaded on-demand and cached -2. **Functional core, imperative shell**: All utility functions are pure; only `EmbedCode` manages state -3. **Interface-based abstraction**: `IInferenceEngine` decouples the model from pure-TypeScript -4. **Self-describing models**: `model-descriptor.json` is the single source of truth for architecture constants — `fromPretrained()` resolves `ModelConfig` via `resolveModelConfig()` from the descriptor -5. **Zero-dependency core for basic usage**: Only `onnxruntime-node` (dynamic import) at inference time -6. **Python parity**: Every source file cites the corresponding HuggingFace/Python source for cross-verification -7. **Progressive disclosure**: Public API exports both high-level (`EmbedCode`) and advanced (`tokenize`, `pool`, `normalize`) APIs - ---- - -## Package Sizes - -| Package | Code Size | Dependencies | -| ---------------------------- | --------- | ---------------- | -| `@agentix-e/embed-code-core` | ~50 KB | embedded weights | -| `@agentix-e/embed-code-cli` | ~10 KB | `commander` | - -Model weights (~137 MB int8 TS for nomic-embed-code 7B) are downloaded separately from GitHub Releases. diff --git a/docs/GETTING-STARTED.md b/docs/GETTING-STARTED.md deleted file mode 100644 index 71c74a8..0000000 --- a/docs/GETTING-STARTED.md +++ /dev/null @@ -1,563 +0,0 @@ -# embed-code-ts Usage Documentation - -> Node.js/TypeScript nomic-embed-code — Code-Aware Text Embeddings - ---- - -## Table of Contents - -1. [Quick Start](#1-quick-start) -2. [Getting the Model](#2-getting-the-model) -3. [API Usage](#3-api-usage) -4. [CLI Tools](#4-cli-tools) -5. [Configuration Reference](#5-configuration-reference) -6. [Output Description](#6-output-description) -7. [Model Update](#7-model-update) -8. [Troubleshooting](#8-troubleshooting) -9. [Performance Guide](#9-performance-guide) - ---- - -## 1. Quick Start - -### Installation - -```bash -git clone https://github.com/AgentiX-E/embed-code-ts.git -cd embed-code-ts -npm install -npm run build:all -``` - -### Minimal Example (requires weights file, see Section 2) - -```typescript -import { EmbedCode, downloadModel } from '@agentix-e/embed-code-core'; - -// 1. Download model (first time only, cached thereafter) -const modelPath = await downloadModel({ - onProgress: (received, total, speed) => { - console.log(`${received.toFixed(0)} / ${total.toFixed(0)} MB @ ${speed.toFixed(1)} MB/s`); - }, -}); - -// 2. Create embedder -const embedder = await EmbedCode.fromPretrained({ modelPath }); - -// 3. Generate embeddings -const results = await embedder.embed([ - 'search_query: How to sort an array in Python?', - 'search_document: def quicksort(arr): return arr if len(arr) <= 1 else quicksort([x for x in arr[1:] if x <= arr[0]]) + [arr[0]] + quicksort([x for x in arr[1:] if x > arr[0]])', -]); - -// 4. Use results -console.log('Embeddings shape:', results.embeddings.length); // 7168 (= 2 × 3584) -console.log('Time:', results.elapsedMs, 'ms'); - -// 5. Compute similarity -const dim = 3584; -const sim = embedder.similarity(results.embeddings.slice(0, dim), results.embeddings.slice(dim)); -console.log('Similarity:', sim); - -// 6. Release resources -await embedder.dispose(); -``` - -> **For full API, configuration, CLI usage** → see subsequent sections of this document. -> **For model export and automation pipeline** → see [MODEL-UPDATE.md](MODEL-UPDATE.md). - ---- - -## 2. Getting the Model - -### 2.1 Auto Download (Recommended) - -```bash - -``` - -```typescript -import { downloadModel } from '@agentix-e/embed-code-core'; - -const modelPath = await downloadModel(); -// → ~/.cache/embed-code-ts/nomic-embed-text-v1.5-int8.weights.bin -``` - -The model is downloaded once, SHA-256 verified, and cached for all subsequent uses. - -### 2.2 One-Click Full Pipeline - -```bash -# Node.js — One-click: Download → Export → Validate → Test → Benchmark -npm run pipeline - -# Or pure Node.js -node scripts/pipeline.js -``` - -> See [MODEL-UPDATE.md](MODEL-UPDATE.md) for details. - -### 2.3 Manual Export from HuggingFace - -**Step 1: Install Python dependencies** - -```bash -pip install torch transformers -``` - -**Step 2: Run the export script** - -```bash -# Basic usage — auto-download nomic-embed-code and export int8 weights -python3 scripts/export-weights.py \ - --output models/nomic-embed-text-v1.5-int8.weights.bin \ - --model nomic-ai/nomic-embed-code \ - --precision int8 - -# Skip validation (faster) -python3 scripts/export-weights.py \ - --output models/nomic-embed-text-v1.5-int8.weights.bin \ - --skip-validation -``` - -**Export flow**: - -``` -HuggingFace Hub Local Disk -┌──────────────────┐ ┌─────────────────────────────────┐ -│ nomic-ai/ │ pip │ models/ │ -│ nomic-embed-code │──────→│ nomic-embed-text-v1.5-int8.weights.bin │ -│ (safetensors) │ export │ (~7 GB int8) │ -└──────────────────┘ └─────────────────────────────────┘ - ~28 GB (fp32) ~7 GB (int8) -``` - -### 2.4 Model File Specifications - -| Property | Value | -| -------- | --------------------------------------------------- | -| Filename | `nomic-embed-code-v1-int8.weights.bin` | -| Size | ~7 GB | -| Format | int8 binary format | -| Input | `input_ids: [batch, 32768]` (int64) | -| Input | `attention_mask: [batch, 32768]` (int64) | -| Output | `last_hidden_state: [batch, 32768, 3584]` (float32) | -| Backend | pure-TypeScript engine | - -### 2.5 Hardware Requirements - -| Component | Minimum | Recommended | -| -------------- | ------------ | ----------- | -| RAM | 512 MB | 2 GB+ | -| Disk | 150 MB free | SSD | -| GPU (Optional) | 2 GB VRAM | 4 GB+ VRAM | -| CPU Mode | ✅ Available | Fast enough | - ---- - -## 3. API Usage - -### 3.1 EmbedCode - -Core class through which all operations are performed. - -```typescript -import { EmbedCode } from '@agentix-e/embed-code-core'; -``` - -#### `fromPretrained(options)` - -Load a pretrained model. - -```typescript -const embedder = await EmbedCode.fromPretrained({ - modelPath: './models/nomic-embed-text-v1.5-int8.weights.bin', // Required - executionProvider: 'cpu', // Optional: 'cpu' | 'cuda' | 'dml' - intraOpNumThreads: 4, // Optional: CPU thread count - skipWarmup: false, // Optional: skip warmup inference - tokenizerPath: '/path/to/tokenizer.json', // Optional: custom tokenizer -}); -``` - -#### `embed(texts, options?)` - -Generate embeddings for one or more texts. - -```typescript -const result = await embedder.embed(texts, { - maxTokens: 32768, - poolingStrategy: 'last_token', - normalize: true, - signal: abortController.signal, - onProgress: (progress) => {}, -}); -// → { embeddings: Float32Array, shape: [N, dim], elapsedMs: number } -``` - -#### `similarity(a, b)` - -Compute cosine similarity between two embedding vectors. - -```typescript -const sim = embedder.similarity(embeddingA, embeddingB); -// → number in range [-1, 1] -``` - -#### `dispose()` - -Release engine resources and GPU memory. - -```typescript -await embedder.dispose(); -``` - -### 3.2 Full Example: Code Search - -```typescript -import { EmbedCode, downloadModel } from '@agentix-e/embed-code-core'; - -async function codeSearch() { - // Load model - const modelPath = await downloadModel(); - const embedder = await EmbedCode.fromPretrained({ modelPath }); - - const { query, document } = embedder.taskPrefixes; - - // Code snippets to search over - const snippets = [ - 'function fibonacci(n) { return n <= 1 ? n : fibonacci(n-1) + fibonacci(n-2); }', - 'function quickSort(arr) { if (arr.length <= 1) return arr; const pivot = arr[0]; return [...quickSort(arr.filter(x => x < pivot)), pivot, ...quickSort(arr.filter(x => x >= pivot))]; }', - 'class BinarySearchTree { insert(value) { /* ... */ } find(value) { /* ... */ } }', - ]; - - // Generate embeddings - const results = await embedder.embed([ - query + 'recursive fibonacci sequence', - ...snippets.map((s) => document + s), - ]); - - const dim = 3584; - const queryEmb = results.embeddings.slice(0, dim); - const snippetEmbs = snippets.map((_, i) => - results.embeddings.slice((i + 1) * dim, (i + 2) * dim), - ); - - // Compute similarities - const similarities = snippetEmbs.map((emb) => embedder.similarity(queryEmb, emb)); - const bestIdx = similarities.indexOf(Math.max(...similarities)); - - console.log('Best match:', snippets[bestIdx]); - console.log('Similarity:', similarities[bestIdx].toFixed(4)); - - await embedder.dispose(); -} - -codeSearch(); -``` - -### 3.3 Using Task Prefixes - -nomic-embed-code requires task prefixes for optimal embedding quality: - -```typescript -const { query, document } = embedder.taskPrefixes; - -// Search query (what you're looking for) -const queryEmbedding = await embedder.embed([query + 'authentication and user login service']); - -// Code document (what you're searching over) -const codeEmbedding = await embedder.embed([document + 'export class AuthService { /* ... */ }']); -``` - ---- - -## 4. CLI Tools - -### Basic Usage - -```bash -# Download model -embed-code setup - -# Generate embedding for text -embed-code embed "function hello() { console.log('Hello World'); }" - -# Embed from a file -embed-code embed -f source.py - -# With custom model -embed-code embed -m ./custom.weights.bin -f source.ts - -# Show model info -embed-code info -embed-code info -m ./custom.weights.bin -``` - -### CLI Parameters - -| Parameter | Required | Description | -| ---------------------- | -------- | ---------------------------------------------------- | -| `` | ❌ | Text to embed (alternative to -f) | -| `-f, --file ` | ❌ | Input file path | -| `-m, --model ` | ❌ | weights file path | -| `-o, --output ` | ❌ | Output file path | -| `--max-tokens ` | ❌ | Max input tokens (default: 32768) | -| `--no-normalize` | ❌ | Disable L2 normalization | -| `--pooling ` | ❌ | Pooling: last_token, mean, cls (default: last_token) | - ---- - -## 5. Configuration Reference - -### Embed Options Full Parameters - -```typescript -interface EmbedOptions { - maxTokens: number; // Max input token count (default: 32768) - poolingStrategy: 'last_token' | 'mean' | 'cls'; // Pooling (default: 'last_token') - normalize: boolean; // L2 normalize output (default: true) - signal?: AbortSignal; // Abort controller signal - onProgress?: (progress: EmbedProgress) => void; // Progress callback -} - -interface EmbedProgress { - phase: 'tokenize' | 'inference' | 'pool' | 'normalize'; - step: number; - total: number; -} -``` - -### Recommended Configuration - -```typescript -// Production -{ - maxTokens: 32768, - poolingStrategy: 'last_token', - normalize: true, -} - -// Maximum accuracy -{ - maxTokens: 32768, - poolingStrategy: 'mean', - normalize: true, -} -``` - -### Input Length Recommendations - -| Scenario | maxTokens | Notes | -| --------------------- | --------- | ---------------------------- | -| Single function/class | 128-256 | Fast, covers most use cases | -| Full source file | 32768 | nomic-embed-code max context | -| Multi-file contexts | N/A | Chunk and embed individually | - ---- - -## 6. Output Description - -### Single Text Embedding - -``` -embeddings: Float32Array(3584) — L2-normalized embedding vector -elapsedMs: number — Inference time -``` - -### Batch Embedding - -``` -embeddings: Float32Array(N × 3584) — Flat array of N embeddings concatenated -elapsedMs: number — Total inference time -``` - -### Embedding Properties - -| Property | Value | Description | -| ---------- | ------------ | ------------------------------------------ | -| Dimensions | 3584 | Fixed embedding size (nomic-embed-code v1) | -| Normalized | true | L2 norm ≈ 1.0 when normalize is enabled | -| Type | Float32Array | 32-bit floating point | -| Range | [-1, 1] | Values typically in [-0.1, 0.1] | - ---- - -## 7. Model Update - -### Check Current nomic-embed-code Latest Version - -nomic-embed-code model versions are published on HuggingFace: -https://huggingface.co/nomic-ai/nomic-embed-code - -### Update to Latest Model - -```bash -# Re-export latest version -python3 scripts/export-weights.py \ - --model nomic-ai/nomic-embed-code \ - --output models/nomic-embed-text-v1.5-int8.weights.bin \ - --precision int8 \ - --skip-validation - -# Verify -npm test -``` - -### Verify Model Compatibility - -```bash -# Run all tests to confirm compatibility -npm test -``` - -If all tests pass, the new model is compatible with current code. - ---- - -## 8. Troubleshooting - -### Model Load Failure - -``` -Error: Engine not loaded. Call load() first. -``` - -**Solution**: - -1. Verify model file exists and is ≥ 100 MB -2. Run `python3 scripts/export-weights.py --output ` to re-export -3. Check that `` is installed - -### Model Download Fails Behind Proxy - -``` -Error: fetch failed -``` - -**Solution**: - -```bash -# Option A: Standard proxy env vars -export HTTPS_PROXY=http://proxy.company.com:8080 - -# Option B: Embed-Code specific env vars -export EMBED_CODE_PROXY_URL=http://proxy.company.com:8080 -export EMBED_CODE_PROXY_USERNAME=user -export EMBED_CODE_PROXY_PASSWORD=pass -``` - -### Out of Memory (OOM) - -``` -JavaScript heap out of memory -``` - -**Solution**: - -```typescript -// Process large batches in chunks -const CHUNK = 50; -const allEmbeddings = []; -for (let i = 0; i < texts.length; i += CHUNK) { - const chunk = texts.slice(i, i + CHUNK); - const { embeddings } = await embedder.embed(chunk); - allEmbeddings.push(embeddings); -} -``` - -### Weights Not Found - -``` -Cannot find module '' -``` - -**Solution**: - -```bash -npm install -``` - ---- - -## 9. Performance Guide - -### Expected Inference Speed - -| Hardware | 1 text (512 tokens) | 32 texts (batch) | -| -------------- | ------------------- | ---------------- | -| CPU (32 cores) | 2-5 ms | 30-60 ms | -| CPU (8 cores) | 5-10 ms | 80-150 ms | -| GPU (8GB) | 0.5-1 ms | 10-20 ms | -| GPU (24GB) | 0.2-0.5 ms | 5-10 ms | - -### Optimization Suggestions - -1. **Batch inference**: Process multiple texts together — 5-10x throughput improvement -2. **Use appropriate maxTokens**: Set to the minimum needed for your use case -3. **GPU acceleration**: `executionProvider: 'cuda'` for 5-20x speedup -4. **Cache the embedder**: Reuse the same `EmbedCode` instance for multiple calls -5. **Pooling strategy**: `last_token` is fastest; `mean` and `cls` are slightly slower but may yield better quality - -### Memory Optimization - -```typescript -// Process large batches in chunks -const CHUNK = 50; -for (let i = 0; i < allTexts.length; i += CHUNK) { - const chunk = allTexts.slice(i, i + CHUNK); - const { embeddings } = await embedder.embed(chunk); - // Save chunk results... -} -``` - ---- - -## Architecture Overview - -``` -┌─────────────────────────────────────────┐ -│ embed-code-ts │ -├─────────────────────────────────────────┤ -│ │ -│ Input Text (string) │ -│ │ │ -│ ▼ │ -│ ┌──────────┐ ┌───────────────┐ │ -│ │Tokenizer │───→│EmbedCodeInfer │ │ -│ │ │ │ pure-TypeScript │ │ -│ │•BPE │ │ │ │ -│ │•Pad/Trunc│ │ C++ native inf│ │ -│ │•Attn Mask│ │ CPU/CUDA/DML │ │ -│ └──────────┘ │ INT8 quantized│ │ -│ └───────────────┘ │ -│ │ │ -│ ▼ │ -│ ┌──────────┐ │ -│ │ Pooler │ │ -│ │ │ │ -│ │•Last Token│ │ -│ │•Mean Pool│ │ -│ │•CLS Pool │ │ -│ └──────────┘ │ -│ │ │ -│ ▼ │ -│ ┌──────────┐ │ -│ │Normalizer│ │ -│ │ │ │ -│ │•L2 Norm │ │ -│ └──────────┘ │ -│ │ │ -│ ▼ │ -│ Output: Float32Array(3584) │ -│ │ -└─────────────────────────────────────────┘ -``` - ---- - -## References - -- **Project Repository**: https://github.com/AgentiX-E/embed-code-ts -- **nomic-embed-code on HuggingFace**: https://huggingface.co/nomic-ai/nomic-embed-code -- **Nomic Embed Paper**: https://arxiv.org/abs/2402.01613 -- **pure-TypeScript engine**: https://github.com/AgentiX-E/embed-code-ts diff --git a/docs/MODEL-UPDATE.md b/docs/MODEL-UPDATE.md deleted file mode 100644 index 8adc622..0000000 --- a/docs/MODEL-UPDATE.md +++ /dev/null @@ -1,166 +0,0 @@ -# nomic-embed-code Model Update Guide - -> This document describes the automated and manual model update flows for embed-code-ts. - ---- - -## Automated Model Updates (Default) - -The project uses a **dual-channel release architecture**: - -### Code Channel - -``` -git tag v* → release.yml - ├─ quality (lint + test) - ├─ publish-npm (OIDC + provenance) - └─ github-release (TypeDoc + release notes) -``` - -Code releases publish npm packages only. They do **not** include weight files. - -### Model Channel - -``` -nightly.yml (cron: 2 AM UTC daily) - │ - ├─ Compare HF revision of nomic-ai/nomic-embed-code - │ against committed models/model-descriptor.json - │ - ├─ New revision detected? → Trigger model-release.yml - │ ├─ detect (idempotency check via model- tag) - │ ├─ export-model (PyTorch → int8 weights with validation) - │ ├─ validate (full test suite + benchmark) - │ ├─ github-release (model- + model-latest tags) - │ └─ update-manifest (commit descriptor → auto-close issue) - │ - └─ No change → no-op -``` - -### Download Channel - -``` -npm install @agentix-e/embed-code-core -node -e "const {downloadModel}=require('@agentix-e/embed-code-core');downloadModel()" - │ - └─ Downloads from: github.com/.../releases/download/model-latest/nomic-embed-text-v1.5-int8.weights.bin - + model-descriptor.json for SHA-256 verification -``` - -The `model-latest` tag is a rolling pointer that always points to the most recently validated model release. Users always get the latest model without upgrading their npm package. - ---- - -## Incbin-Inspired Architecture - -The incbin-inspired approach keeps the npm package code-only while delivering models via on-demand download: - -``` -npm package (@agentix-e/embed-code-core) -│ ~50 KB TypeScript/JS code only -│ + model-descriptor.json (SHA-256 fingerprint) -│ -└─ On first use: downloadModel() - ├─ Fetches int8 weights from GitHub Releases - ├─ SHA-256 verifies against descriptor - └─ Caches to ~/.cache/embed-code-ts/ - -Subsequent runs: pure filesystem read — zero network, zero latency -``` - ---- - -## Manual Model Update - -If you need to manually trigger a model export and release: - -### 1. Trigger model release workflow - -```bash -# Automatic (force re-release of current HF revision) -gh workflow run model-release.yml -f force=true - -# Or via GitHub UI: Actions → Model Release → Run workflow -``` - -### 2. Local export for testing - -```bash -# Full pipeline: check HF → export → validate → test → benchmark -npm run pipeline - -# Export only (skip tests) -npm run pipeline:export -``` - -### 3. Update the committed descriptor - -```bash -# After local export with a new HF revision -git add models/model-descriptor.json -git commit -m "chore(model): update descriptor to HF rev " -``` - ---- - -## ModelDescriptor Contract - -The `model-descriptor.json` file (committed to the repo and distributed with each weight file release) defines the architecture contract between the model and the TypeScript engine: - -| Field | Source | Purpose | -| ------------------- | ------------------------- | ------------------------------------------ | -| `schema` | Constant (1) | Forward compatibility version | -| `model.hf_revision` | HuggingFace API | Traceability to exact PyTorch checkpoint | -| `weights` | weight file | Runtime shape validation | -| `weights.sha256` | Computed from weight file | Download integrity verification | -| `architecture.*` | PyTorch model params | Configures the TypeScript engine | -| `tokenizer.*` | Model config | Vocabulary size, maxTokens, special tokens | - -The engine reads this descriptor at runtime via `loadModelDescriptor()` and converts it to a `ModelConfig` via `descriptorToModelConfig()`. All hardcoded architecture constants in the TypeScript code have been eliminated — the descriptor is the single source of truth. - ---- - -## Version Compatibility - -| Schema | Engine Requirement | Notes | -| ------ | ------------------------------------ | --------------------------------------------------------------- | -| 1 | `@agentix-e/embed-code-core` ≥ 0.1.0 | Current | -| > 1 | Upgrade required | Engine logs warning and falls back to `NOMIC_EMBED_CODE_CONFIG` | - -If a future nomic-embed-code model version has a different architecture, updating `export-weights.py` to generate the new descriptor is sufficient — no TypeScript code changes are needed as long as the schema version is compatible. - ---- - -## Model Update Flow - -When a new model version is released: - -``` -1. New HF model detected - │ -2. Export int8 weights (~137 MB) - │ -3. Validate: - ├─ Model loads correctly ✓ - ├─ Embedding dimensions match ✓ - ├─ Full test suite passes ✓ - └─ Benchmark regression ≤ 5% ✓ - │ -4. Release: - ├─ GitHub Release: publish int8 weights + descriptor - └─ Update model-latest tag -``` - ---- - -## Troubleshooting - -| Issue | Solution | -| ------------------------------------ | ------------------------------------------------------------------------------------------------- | -| Nightly check not detecting changes | Verify `models/model-descriptor.json` is committed with the current HF revision | -| Model release fails | Check the workflow run logs; force re-run with `force: true` | -| `downloadModel()` fails | Verify `model-latest` release exists and contains both `.weights.bin` and `model-descriptor.json` | -| `downloadModel()` fails behind proxy | Set `HTTPS_PROXY` or `EMBED_CODE_PROXY_URL` environment variables (see README Proxy section) | -| `downloadModel()` fails with 407 | Proxy authentication required — set `EMBED_CODE_PROXY_USERNAME` and `EMBED_CODE_PROXY_PASSWORD` | -| Engine rejects model | Descriptor schema > `ENGINE_SUPPORTED_SCHEMA` — upgrade `@agentix-e/embed-code-core` | -| Local tests need weight file | Export locally: `python3 scripts/export-weights.py` or `npm run pipeline:export` | diff --git a/docs/index.html b/docs/index.html deleted file mode 100644 index fc416fc..0000000 --- a/docs/index.html +++ /dev/null @@ -1,72 +0,0 @@ - - - - - embed-code-ts · Docs - - - -

🚀 embed-code-ts

-

- Node.js/TypeScript implementation of Nomic's nomic-embed-code — code-aware text - embeddings with INT8 weights packed directly into the npm package. -

-
-

📚 API Documentation

-

Full TypeDoc reference for all packages

-
-
-

📊 Benchmark Reports

-

- Inference latency, throughput, and embedding quality benchmarks (Node.js + WASM comparison) -

-
-
-

📈 Test Coverage

-

Code coverage reports (lines, branches, functions, statements)

-
-
-

💻 Source Code

-

GitHub repository with README, contributing guide, and full source

-
- - diff --git a/packages/embed-code-core/test/unit/batch-processor.test.ts b/packages/embed-code-core/test/unit/batch-processor.test.ts index 09daaf2..7bb6dd8 100644 --- a/packages/embed-code-core/test/unit/batch-processor.test.ts +++ b/packages/embed-code-core/test/unit/batch-processor.test.ts @@ -65,4 +65,35 @@ describe('processBatch', () => { }); expect(results.length).toBe(2); }); + + it('uses default concurrency from os.cpus()', async () => { + const results: number[] = []; + // No concurrency option — triggers importNodeCores() path + await processBatch([10, 20, 30], async (n) => { + results.push(n); + }); + expect(new Set(results)).toEqual(new Set([10, 20, 30])); + }); + + it('times out on slow items', async () => { + const results: number[] = []; + await processBatch( + [1, 2], + async (n) => { + if (n === 1) await new Promise((r) => setTimeout(r, 200)); + results.push(n); + }, + { concurrency: 2, timeout: 50 }, + ); + // Item 1 should have timed out, item 2 should complete + expect(results).toContain(2); + expect(results).not.toContain(1); + }); + + it('handles items exactly at concurrency boundary', async () => { + const results: number[] = []; + // concurrency=2, items=2: worker loop breaks at i>=2 immediately + await processBatch([5, 6], async (n) => results.push(n), { concurrency: 3 }); + expect(results).toEqual([5, 6]); + }); }); diff --git a/packages/embed-code-core/test/unit/descriptor.test.ts b/packages/embed-code-core/test/unit/descriptor.test.ts index aa44c87..23035bc 100644 --- a/packages/embed-code-core/test/unit/descriptor.test.ts +++ b/packages/embed-code-core/test/unit/descriptor.test.ts @@ -157,6 +157,91 @@ describe('resolveModelConfig', () => { fs.rmSync(tmpDir, { recursive: true, force: true }); }); + it('resolves config with weights field instead of onnx', () => { + const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'embed-test-')); + const modelPath = path.join(tmpDir, 'model.weights.bin'); + fs.writeFileSync(modelPath, 'dummy content'); + + const descriptor = { + schema: 1, + model: { + name: 'nomic-embed-code', + version: 'v1', + base_architecture: 'BERT-base', + hf_repository: 'nomic-ai/nomic-embed-code', + hf_revision: 'abc123', + exported_at: '2026-01-01T00:00:00Z', + precision: 'int8', + }, + weights: { + input_ids_name: 'custom_input_ids', + attention_mask_name: 'custom_attention_mask', + output_name: 'custom_output', + }, + architecture: { + embedding_dim: 768, + num_layers: 12, + num_heads: 12, + num_kv_heads: 0, + head_dim: 64, + hidden_size: 768, + intermediate_size: 3072, + vocab_size: 40856, + max_position_embeddings: 8192, + rope_theta: 0, + sliding_window: null, + attention_dropout: 0.0, + use_sliding_window: false, + }, + tokenizer: { + type: 'bpe', + vocab_size: 40856, + max_length: 8192, + pad_token: '<|endoftext|>', + pad_token_id: 0, + bos_token: null, + bos_token_id: null, + eos_token: '<|endoftext|>', + eos_token_id: 0, + unk_token: null, + unk_token_id: null, + }, + pooling: { strategy: 'mean', normalize: true }, + task_prefixes: { query: 'search_query: ', document: 'search_document: ' }, + }; + + fs.writeFileSync(path.join(tmpDir, 'model-descriptor.json'), JSON.stringify(descriptor)); + + const { config } = resolveModelConfig(modelPath); + expect(config.inputIdsName).toBe('custom_input_ids'); + expect(config.attentionMaskName).toBe('custom_attention_mask'); + expect(config.outputName).toBe('custom_output'); + + fs.rmSync(tmpDir, { recursive: true, force: true }); + }); + + it('falls back when descriptor has no weights or onnx', () => { + const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'embed-test-')); + const modelPath = path.join(tmpDir, 'model.weights.bin'); + fs.writeFileSync(modelPath, 'dummy content'); + + const descriptor = { + schema: 1, + model: { name: 'test' }, + architecture: { embedding_dim: 768, max_position_embeddings: 8192 }, + pooling: { strategy: 'mean', normalize: true }, + task_prefixes: { query: 'q', document: 'd' }, + }; + + fs.writeFileSync(path.join(tmpDir, 'model-descriptor.json'), JSON.stringify(descriptor)); + + const { config } = resolveModelConfig(modelPath); + expect(config.inputIdsName).toBe('input_ids'); // fallback + expect(config.attentionMaskName).toBe('attention_mask'); + expect(config.outputName).toBe('last_hidden_state'); + + fs.rmSync(tmpDir, { recursive: true, force: true }); + }); it('uses explicit fallback config when provided', () => { const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'embed-test-')); const modelPath = path.join(tmpDir, 'model.weights.bin'); diff --git a/packages/embed-code-core/test/unit/pooling.test.ts b/packages/embed-code-core/test/unit/pooling.test.ts index 977a6e9..23dfebd 100644 --- a/packages/embed-code-core/test/unit/pooling.test.ts +++ b/packages/embed-code-core/test/unit/pooling.test.ts @@ -78,4 +78,14 @@ describe('cosineSimilarity', () => { it('dimension mismatch produces 0', () => { expect(cosineSimilarity([1, 2], [1, 2, 3])).toBe(0); }); + + it('zero vectors produce 0', () => { + expect(cosineSimilarity([0, 0], [0, 0])).toBe(0); + }); + + it('handles Float32Array inputs', () => { + const a = new Float32Array([1, 0]); + const b = new Float32Array([0, 1]); + expect(cosineSimilarity(a, b)).toBeCloseTo(0.0); + }); }); diff --git a/packages/embed-code-core/test/unit/pre-tokenizer.test.ts b/packages/embed-code-core/test/unit/pre-tokenizer.test.ts index 3911bf7..46c3107 100644 --- a/packages/embed-code-core/test/unit/pre-tokenizer.test.ts +++ b/packages/embed-code-core/test/unit/pre-tokenizer.test.ts @@ -45,8 +45,28 @@ describe('isPunctuation', () => { expect(isPunctuation(',')).toBe(true); }); + it('range 33-47: exclamation mark', () => { + expect(isPunctuation('!')).toBe(true); + }); + + it('range 58-64: colon', () => { + expect(isPunctuation(':')).toBe(true); + }); + + it('range 91-96: backslash', () => { + expect(isPunctuation('\\')).toBe(true); + }); + + it('range 123-126: tilde', () => { + expect(isPunctuation('~')).toBe(true); + }); + it('returns false for letters', () => { expect(isPunctuation('a')).toBe(false); expect(isPunctuation('Z')).toBe(false); }); + + it('returns false for digits', () => { + expect(isPunctuation('5')).toBe(false); + }); }); diff --git a/packages/embed-code-core/test/unit/tokenizer.test.ts b/packages/embed-code-core/test/unit/tokenizer.test.ts index 0e3d9b9..3da0648 100644 --- a/packages/embed-code-core/test/unit/tokenizer.test.ts +++ b/packages/embed-code-core/test/unit/tokenizer.test.ts @@ -37,6 +37,17 @@ describe('WordPieceTokenizer', () => { }); }); + describe('static constructors', () => { + it('fromBuffer works with ArrayBuffer', () => { + if (!hasTokenizer) return; + const raw = fs.readFileSync(TOKENIZER_PATH); + const buf = raw.buffer.slice(raw.byteOffset, raw.byteOffset + raw.byteLength); + const tok = WordPieceTokenizer.fromBuffer(buf, 512); + expect(tok.clsTokenId).toBe(101); + expect(tok.sepTokenId).toBe(102); + }); + }); + describe('basic tokenization', () => { let tok: WordPieceTokenizer; diff --git a/packages/embed-code-core/test/unit/vocab-loader.test.ts b/packages/embed-code-core/test/unit/vocab-loader.test.ts index bb7ba7c..f0ff143 100644 --- a/packages/embed-code-core/test/unit/vocab-loader.test.ts +++ b/packages/embed-code-core/test/unit/vocab-loader.test.ts @@ -29,4 +29,24 @@ describe('loadVocab', () => { const info = loadVocab(json); expect(info.padTokenId).toBe(0); // default fallback }); + + it('recognizes continuation subwords from added_tokens', () => { + const json = { + model: { vocab: { hello: 1, world: 2 } }, + added_tokens: [ + { content: '##ing', id: 3, special: false }, + { content: '##ly', id: 4, special: false }, + ], + }; + const info = loadVocab(json); + expect(info.isContinuation(3)).toBe(true); + expect(info.isContinuation(4)).toBe(true); + expect(info.isContinuation(1)).toBe(false); + }); + + it('handles empty model.vocab', () => { + const info = loadVocab({ model: { vocab: {} } }); + expect(info.size).toBe(0); + expect(info.padTokenId).toBe(0); // fallback to 0 + }); }); diff --git a/packages/embed-code-node/test/onnx-embedder.test.ts b/packages/embed-code-node/test/onnx-embedder.test.ts new file mode 100644 index 0000000..5547db4 --- /dev/null +++ b/packages/embed-code-node/test/onnx-embedder.test.ts @@ -0,0 +1,75 @@ +import { describe, it, expect } from 'vitest'; +import { NodeEmbedder } from '../src/onnx-embedder'; +import * as fs from 'node:fs'; +import * as path from 'node:path'; + +const MODEL_PATH = process.env.EMBED_CODE_MODEL_PATH || 'models/nomic-embed-code-v1.5.int8.onnx'; +const hasModel = fs.existsSync(MODEL_PATH); + +describe('NodeEmbedder', () => { + it('create fails when tokenizer missing', async () => { + await expect(NodeEmbedder.create({ modelPath: '/nonexistent/model.onnx' })).rejects.toThrow( + 'Tokenizer not found', + ); + }); + + it('create fails when model missing', async () => { + const dir = path.join(process.cwd(), 'models'); + await expect( + NodeEmbedder.create({ modelPath: path.join(dir, 'nonexistent.onnx') }), + ).rejects.toThrow(); + }); + + it('embed with real model returns 768-dim normalized vector', async () => { + if (!hasModel) return; + const embedder = await NodeEmbedder.create({ modelPath: MODEL_PATH }); + const result = await embedder.embed('search_query: hello world'); + expect(result).toBeInstanceOf(Float32Array); + expect(result.length).toBe(768); + // L2-normalized: norm should be ~1.0 + let norm = 0; + for (let i = 0; i < result.length; i++) norm += result[i] ** 2; + expect(Math.sqrt(norm)).toBeCloseTo(1, 1); + await embedder.dispose(); + }, 120000); + + it('embedBatch with real model returns array of embeddings', async () => { + if (!hasModel) return; + const embedder = await NodeEmbedder.create({ modelPath: MODEL_PATH }); + const results = await embedder.embedBatch(['query one', 'query two']); + expect(results.length).toBe(2); + expect(results[0].length).toBe(768); + expect(results[1].length).toBe(768); + await embedder.dispose(); + }, 120000); + + it('embedBatch with onProgress callback', async () => { + if (!hasModel) return; + const embedder = await NodeEmbedder.create({ modelPath: MODEL_PATH }); + const progress: number[] = []; + const results = await embedder.embedBatch(['a', 'b', 'c'], { + concurrency: 2, + onProgress: (done, _total) => progress.push(done), + }); + expect(results.length).toBe(3); + expect(progress.length).toBeGreaterThanOrEqual(3); + await embedder.dispose(); + }, 120000); + + it('embed throws when session disposed', async () => { + if (!hasModel) return; + const embedder = await NodeEmbedder.create({ modelPath: MODEL_PATH }); + await embedder.dispose(); + await expect(embedder.embed('test')).rejects.toThrow('Session not initialized'); + }); + + it('modelInfo reflects correct dimensions', async () => { + if (!hasModel) return; + const embedder = await NodeEmbedder.create({ modelPath: MODEL_PATH }); + expect(embedder.dimensions).toBe(768); + expect(embedder.maxSequenceLength).toBe(512); + expect(embedder.modelInfo.name).toBe('nomic-embed-code'); + expect(embedder.modelInfo.quantization).toBe('int8'); + await embedder.dispose(); + }); +}); diff --git a/packages/embed-code-node/test/ort-backend.test.ts b/packages/embed-code-node/test/ort-backend.test.ts new file mode 100644 index 0000000..c8b8bf8 --- /dev/null +++ b/packages/embed-code-node/test/ort-backend.test.ts @@ -0,0 +1,38 @@ +import { describe, it, expect } from 'vitest'; +import { NodeOrtBackend } from '../src/ort-backend'; + +describe('NodeOrtBackend', () => { + const backend = new NodeOrtBackend(); + + it('createTensor for int64 from Int32Array', () => { + const tensor = backend.createTensor('int64', new Int32Array([1, 2, 3]), [3]); + expect(tensor.type).toBe('int64'); + expect(tensor.dims).toEqual([3]); + expect(tensor.data).toBeInstanceOf(BigInt64Array); + }); + + it('createTensor for int64 from number[]', () => { + const tensor = backend.createTensor('int64', [1, 2, 3], [3]); + expect(tensor.type).toBe('int64'); + expect(tensor.dims).toEqual([3]); + expect(tensor.data).toBeInstanceOf(BigInt64Array); + }); + + it('createTensor for float32 from Float32Array', () => { + const tensor = backend.createTensor('float32', new Float32Array([1.0, 2.0, 3.0]), [3]); + expect(tensor.type).toBe('float32'); + expect(tensor.dims).toEqual([3]); + expect(tensor.data).toBeInstanceOf(Float32Array); + }); + + it('createTensor for float32 from number[]', () => { + const tensor = backend.createTensor('float32', [1.0, 2.0], [2]); + expect(tensor.type).toBe('float32'); + expect(tensor.dims).toEqual([2]); + }); + + it('dispose tensor does nothing', () => { + const tensor = backend.createTensor('float32', [1.0, 2.0], [2]); + expect(() => tensor.dispose()).not.toThrow(); + }); +}); diff --git a/packages/embed-code-web/src/onnx-embedder.ts b/packages/embed-code-web/src/onnx-embedder.ts index b5dac2e..9098c9e 100644 --- a/packages/embed-code-web/src/onnx-embedder.ts +++ b/packages/embed-code-web/src/onnx-embedder.ts @@ -39,8 +39,15 @@ export class WebEmbedder implements IEmbedder { static async create(modelUrl: string, tokenizerJson: Record): Promise { const tokenizer = WordPieceTokenizer.fromJSON(tokenizerJson, 512); - const resp = await fetch(modelUrl); - const buffer = await resp.arrayBuffer(); + let buffer: ArrayBuffer; + if (modelUrl.startsWith('file://')) { + const fs = await import('node:fs'); + const raw = fs.readFileSync(modelUrl.replace('file://', '')); + buffer = raw.buffer.slice(raw.byteOffset, raw.byteOffset + raw.byteLength); + } else { + const resp = await fetch(modelUrl); + buffer = await resp.arrayBuffer(); + } const embedder = new WebEmbedder(tokenizer, buffer); const ort = await getOrt(); diff --git a/packages/embed-code-web/test/onnx-embedder.test.ts b/packages/embed-code-web/test/onnx-embedder.test.ts new file mode 100644 index 0000000..0193f26 --- /dev/null +++ b/packages/embed-code-web/test/onnx-embedder.test.ts @@ -0,0 +1,81 @@ +import { describe, it, expect } from 'vitest'; +import { WebEmbedder } from '../src/onnx-embedder'; +import * as fs from 'node:fs'; +import * as path from 'node:path'; + +const MODEL_PATH = process.env.EMBED_CODE_MODEL_PATH || 'models/nomic-embed-code-v1.5.int8.onnx'; +const hasModel = fs.existsSync(MODEL_PATH); + +describe('WebEmbedder', () => { + it('embed with real model via WASM returns 768-dim normalized vector', async () => { + if (!hasModel) return; + const tokenizerJson = JSON.parse( + fs.readFileSync(path.join(path.dirname(MODEL_PATH), 'tokenizer.json'), 'utf-8'), + ); + const modelUrl = `file://${path.resolve(MODEL_PATH)}`; + const embedder = await WebEmbedder.create(modelUrl, tokenizerJson); + const result = await embedder.embed('search_query: hello world'); + expect(result).toBeInstanceOf(Float32Array); + expect(result.length).toBe(768); + // L2-normalized: norm should be ~1.0 + let norm = 0; + for (let i = 0; i < result.length; i++) norm += result[i] ** 2; + expect(Math.sqrt(norm)).toBeCloseTo(1, 1); + await embedder.dispose(); + }, 120000); + + it('embedBatch with real model via WASM returns array of embeddings', async () => { + if (!hasModel) return; + const tokenizerJson = JSON.parse( + fs.readFileSync(path.join(path.dirname(MODEL_PATH), 'tokenizer.json'), 'utf-8'), + ); + const modelUrl = `file://${path.resolve(MODEL_PATH)}`; + const embedder = await WebEmbedder.create(modelUrl, tokenizerJson); + const results = await embedder.embedBatch(['query one', 'query two']); + expect(results.length).toBe(2); + expect(results[0].length).toBe(768); + expect(results[1].length).toBe(768); + await embedder.dispose(); + }, 120000); + + it('embedBatch with onProgress callback', async () => { + if (!hasModel) return; + const tokenizerJson = JSON.parse( + fs.readFileSync(path.join(path.dirname(MODEL_PATH), 'tokenizer.json'), 'utf-8'), + ); + const modelUrl = `file://${path.resolve(MODEL_PATH)}`; + const embedder = await WebEmbedder.create(modelUrl, tokenizerJson); + const progress: number[] = []; + const results = await embedder.embedBatch(['a', 'b', 'c'], { + onProgress: (done, _total) => progress.push(done), + }); + expect(results.length).toBe(3); + expect(progress.length).toBeGreaterThanOrEqual(3); + await embedder.dispose(); + }, 120000); + + it('embed throws when session disposed', async () => { + if (!hasModel) return; + const tokenizerJson = JSON.parse( + fs.readFileSync(path.join(path.dirname(MODEL_PATH), 'tokenizer.json'), 'utf-8'), + ); + const modelUrl = `file://${path.resolve(MODEL_PATH)}`; + const embedder = await WebEmbedder.create(modelUrl, tokenizerJson); + await embedder.dispose(); + await expect(embedder.embed('test')).rejects.toThrow('Session not initialized'); + }); + + it('modelInfo reflects correct dimensions', async () => { + if (!hasModel) return; + const tokenizerJson = JSON.parse( + fs.readFileSync(path.join(path.dirname(MODEL_PATH), 'tokenizer.json'), 'utf-8'), + ); + const modelUrl = `file://${path.resolve(MODEL_PATH)}`; + const embedder = await WebEmbedder.create(modelUrl, tokenizerJson); + expect(embedder.dimensions).toBe(768); + expect(embedder.maxSequenceLength).toBe(512); + expect(embedder.modelInfo.name).toBe('nomic-embed-code'); + expect(embedder.modelInfo.quantization).toBe('int8'); + await embedder.dispose(); + }); +}); diff --git a/vitest.config.ts b/vitest.config.ts index 36566ed..69b3484 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -22,16 +22,25 @@ export default defineConfig({ include: [ 'packages/embed-code-core/src/**/*.ts', 'packages/embed-code-node/src/**/*.ts', + 'packages/embed-code-web/src/**/*.ts', 'packages/embed-code-cli/src/**/*.ts', ], exclude: [ 'packages/*/src/index.ts', 'packages/embed-code-cli/src/cli.ts', 'packages/embed-code-core/src/types.ts', + 'packages/embed-code-core/src/embedder-interface.ts', + 'packages/embed-code-core/src/ort-backend-interface.ts', 'packages/*/src/types/**/*.d.ts', ], reporter: ['text', 'html', 'json-summary', 'lcov'], reportsDirectory: './coverage', + thresholds: { + lines: 95, + branches: 90, + functions: 95, + statements: 95, + }, }, }, }); diff --git a/vitest.unit.config.ts b/vitest.unit.config.ts index 9840e9f..4493319 100644 --- a/vitest.unit.config.ts +++ b/vitest.unit.config.ts @@ -29,6 +29,12 @@ export default defineConfig({ ], reporter: ['text', 'text-summary', 'json-summary', 'lcov', 'html'], reportsDirectory: './coverage', + thresholds: { + lines: 95, + branches: 90, + functions: 95, + statements: 95, + }, }, }, });