d803bfe2b1
git-subtree-dir: vendor/ruvector git-subtree-split: b64c21726f2bb37286d9ee36a7869fef60cc6900
1299 lines
39 KiB
JavaScript
1299 lines
39 KiB
JavaScript
/**
|
|
* MicroLoRA - Lightweight LoRA Customization SDK for End-User Model Adaptation
|
|
*
|
|
* Enables browser-based fine-tuning of small LLMs using Low-Rank Adaptation (LoRA).
|
|
* Optimized for edge deployment with minimal memory footprint.
|
|
*
|
|
* @module @ruvector/edge-net/models/microlora
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* import { MicroLoRA } from '@ruvector/edge-net/models';
|
|
*
|
|
* const lora = new MicroLoRA('phi-1.5-int4', { rank: 8 });
|
|
* await lora.train([
|
|
* { input: 'translate to python', output: 'def main():' },
|
|
* { input: 'write a loop', output: 'for i in range(10):' }
|
|
* ], { epochs: 10, lr: 1e-4 });
|
|
*
|
|
* const result = await lora.generate('write a function');
|
|
* console.log(result.text);
|
|
*
|
|
* await lora.saveAdapter('my-code-adapter');
|
|
* ```
|
|
*/
|
|
|
|
import { EventEmitter } from 'events';
|
|
import { createHash, randomBytes } from 'crypto';
|
|
|
|
// ============================================
|
|
// TYPE DEFINITIONS (JSDoc)
|
|
// ============================================
|
|
|
|
/**
|
|
* @typedef {Object} MicroLoRAConfig
|
|
* @property {number} [rank=4] - LoRA rank (dimension of low-rank matrices)
|
|
* @property {number} [alpha=8] - LoRA alpha (scaling factor)
|
|
* @property {number} [dropout=0.05] - Dropout rate during training
|
|
* @property {string[]} [targetModules=['query', 'value']] - Modules to adapt
|
|
* @property {boolean} [quantized=true] - Use quantized weights
|
|
* @property {number} [embeddingDim=384] - Embedding dimension
|
|
*/
|
|
|
|
/**
|
|
* @typedef {Object} TrainingExample
|
|
* @property {string} input - Input text/prompt
|
|
* @property {string} output - Expected output
|
|
* @property {number} [quality=1.0] - Example quality weight
|
|
* @property {Object} [metadata] - Optional metadata
|
|
*/
|
|
|
|
/**
|
|
* @typedef {Object} TrainingOptions
|
|
* @property {number} [epochs=10] - Number of training epochs
|
|
* @property {number} [lr=1e-4] - Learning rate
|
|
* @property {number} [batchSize=4] - Training batch size
|
|
* @property {string} [scheduler='cosine'] - LR scheduler type
|
|
* @property {number} [warmupSteps=10] - Warmup steps
|
|
* @property {boolean} [useEWC=false] - Use EWC for continual learning
|
|
* @property {number} [ewcLambda=1000] - EWC regularization strength
|
|
* @property {boolean} [gradientCheckpointing=true] - Memory optimization
|
|
* @property {Function} [onProgress] - Progress callback
|
|
*/
|
|
|
|
/**
|
|
* @typedef {Object} GenerationOptions
|
|
* @property {number} [maxTokens=64] - Maximum tokens to generate
|
|
* @property {number} [temperature=0.7] - Sampling temperature
|
|
* @property {number} [topP=0.9] - Top-p sampling
|
|
* @property {number} [topK=50] - Top-k sampling
|
|
* @property {number} [repetitionPenalty=1.1] - Repetition penalty
|
|
*/
|
|
|
|
/**
|
|
* @typedef {Object} AdapterMetadata
|
|
* @property {string} id - Unique adapter ID
|
|
* @property {string} name - Human-readable name
|
|
* @property {string} description - Adapter description
|
|
* @property {string} baseModel - Base model used
|
|
* @property {string} domain - Domain category
|
|
* @property {number} rank - LoRA rank
|
|
* @property {number} alpha - LoRA alpha
|
|
* @property {number} trainingSamples - Number of training samples
|
|
* @property {number} trainingEpochs - Training epochs
|
|
* @property {number} createdAt - Creation timestamp
|
|
* @property {string} version - Adapter version
|
|
*/
|
|
|
|
// ============================================
|
|
// CONSTANTS
|
|
// ============================================
|
|
|
|
/**
|
|
* Supported base models for MicroLoRA
|
|
*/
|
|
export const SUPPORTED_MODELS = {
|
|
'phi-1.5-int4': {
|
|
id: 'Xenova/phi-1_5',
|
|
name: 'Phi-1.5 INT4',
|
|
size: '~280MB',
|
|
embeddingDim: 2048,
|
|
hiddenDim: 2048,
|
|
numLayers: 24,
|
|
capabilities: ['code', 'reasoning', 'math'],
|
|
},
|
|
'phi-2-int4': {
|
|
id: 'Xenova/phi-2',
|
|
name: 'Phi-2 INT4',
|
|
size: '~550MB',
|
|
embeddingDim: 2560,
|
|
hiddenDim: 2560,
|
|
numLayers: 32,
|
|
capabilities: ['code', 'reasoning', 'math', 'general'],
|
|
},
|
|
'distilgpt2': {
|
|
id: 'Xenova/distilgpt2',
|
|
name: 'DistilGPT-2',
|
|
size: '~82MB',
|
|
embeddingDim: 768,
|
|
hiddenDim: 768,
|
|
numLayers: 6,
|
|
capabilities: ['general', 'completion'],
|
|
},
|
|
'gpt2': {
|
|
id: 'Xenova/gpt2',
|
|
name: 'GPT-2',
|
|
size: '~250MB',
|
|
embeddingDim: 768,
|
|
hiddenDim: 768,
|
|
numLayers: 12,
|
|
capabilities: ['general', 'completion', 'creative'],
|
|
},
|
|
'starcoder-tiny': {
|
|
id: 'Xenova/tiny_starcoder_py',
|
|
name: 'StarCoder Tiny',
|
|
size: '~40MB',
|
|
embeddingDim: 768,
|
|
hiddenDim: 768,
|
|
numLayers: 6,
|
|
capabilities: ['code', 'python'],
|
|
},
|
|
'qwen-0.5b': {
|
|
id: 'Xenova/Qwen1.5-0.5B',
|
|
name: 'Qwen 0.5B',
|
|
size: '~430MB',
|
|
embeddingDim: 1024,
|
|
hiddenDim: 2816,
|
|
numLayers: 24,
|
|
capabilities: ['multilingual', 'general', 'code'],
|
|
},
|
|
};
|
|
|
|
/**
|
|
* Default MicroLoRA configuration
|
|
*/
|
|
const DEFAULT_CONFIG = {
|
|
rank: 4,
|
|
alpha: 8,
|
|
dropout: 0.05,
|
|
targetModules: ['query', 'value', 'key', 'dense'],
|
|
quantized: true,
|
|
embeddingDim: 384,
|
|
};
|
|
|
|
// ============================================
|
|
// MICROLORA CLASS
|
|
// ============================================
|
|
|
|
/**
|
|
* MicroLoRA - End-user model adaptation SDK
|
|
*
|
|
* Provides a complete workflow for fine-tuning small LLMs in the browser
|
|
* using Low-Rank Adaptation. Optimized for edge computing with support
|
|
* for gradient checkpointing, EWC continual learning, and ONNX export.
|
|
*
|
|
* @extends EventEmitter
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* // Initialize with model and config
|
|
* const lora = new MicroLoRA('phi-1.5-int4', { rank: 8, alpha: 16 });
|
|
*
|
|
* // Train on examples
|
|
* await lora.train([
|
|
* { input: 'Hello', output: 'World' },
|
|
* { input: 'Goodbye', output: 'Friend' },
|
|
* ], { epochs: 5 });
|
|
*
|
|
* // Generate with adapter
|
|
* const result = await lora.generate('Hello there');
|
|
*
|
|
* // Save and share
|
|
* await lora.saveAdapter('./my-adapter.json');
|
|
* ```
|
|
*/
|
|
export class MicroLoRA extends EventEmitter {
|
|
/**
|
|
* Create a MicroLoRA instance
|
|
*
|
|
* @param {string} baseModel - Base model identifier (e.g., 'phi-1.5-int4')
|
|
* @param {MicroLoRAConfig} [config={}] - LoRA configuration
|
|
*/
|
|
constructor(baseModel, config = {}) {
|
|
super();
|
|
|
|
this.id = `microlora-${randomBytes(6).toString('hex')}`;
|
|
this.baseModelKey = baseModel;
|
|
this.baseModel = SUPPORTED_MODELS[baseModel] || {
|
|
id: baseModel,
|
|
name: baseModel,
|
|
embeddingDim: config.embeddingDim || 384,
|
|
hiddenDim: config.embeddingDim || 384,
|
|
numLayers: 12,
|
|
capabilities: ['general'],
|
|
};
|
|
|
|
/** @type {MicroLoRAConfig} */
|
|
this.config = {
|
|
...DEFAULT_CONFIG,
|
|
embeddingDim: this.baseModel.embeddingDim,
|
|
...config,
|
|
};
|
|
|
|
// Initialize adapter weights
|
|
this.adapters = this._initializeAdapters();
|
|
|
|
// Training state
|
|
this.trainingState = null;
|
|
this.isTraining = false;
|
|
|
|
// EWC state for continual learning
|
|
this.ewcState = null;
|
|
|
|
// Inference pipeline
|
|
this.pipeline = null;
|
|
this.initialized = false;
|
|
|
|
// Statistics
|
|
this.stats = {
|
|
totalTrainingSamples: 0,
|
|
totalTrainingTime: 0,
|
|
totalInferences: 0,
|
|
corrections: 0,
|
|
adaptations: 0,
|
|
};
|
|
|
|
// Metadata for adapter
|
|
this.metadata = {
|
|
id: this.id,
|
|
name: 'Untitled Adapter',
|
|
description: '',
|
|
baseModel: this.baseModelKey,
|
|
domain: 'general',
|
|
rank: this.config.rank,
|
|
alpha: this.config.alpha,
|
|
trainingSamples: 0,
|
|
trainingEpochs: 0,
|
|
createdAt: Date.now(),
|
|
version: '1.0.0',
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Initialize LoRA adapter matrices for each target module
|
|
* @private
|
|
* @returns {Map<string, Object>} Adapter weights per module
|
|
*/
|
|
_initializeAdapters() {
|
|
const adapters = new Map();
|
|
const inputDim = this.config.embeddingDim;
|
|
const outputDim = this.config.embeddingDim;
|
|
const rank = this.config.rank;
|
|
|
|
for (const moduleName of this.config.targetModules) {
|
|
adapters.set(moduleName, {
|
|
// A: (inputDim x rank) - initialized with small Gaussian
|
|
loraA: this._kaiming(inputDim, rank),
|
|
// B: (rank x outputDim) - initialized to zero
|
|
loraB: this._zeros(rank, outputDim),
|
|
// Scaling factor
|
|
scaling: this.config.alpha / this.config.rank,
|
|
});
|
|
}
|
|
|
|
return adapters;
|
|
}
|
|
|
|
/**
|
|
* Kaiming/He initialization for matrix
|
|
* @private
|
|
*/
|
|
_kaiming(rows, cols) {
|
|
const matrix = [];
|
|
const std = Math.sqrt(2 / (rows + cols));
|
|
for (let i = 0; i < rows; i++) {
|
|
const row = [];
|
|
for (let j = 0; j < cols; j++) {
|
|
row.push(this._gaussianRandom() * std);
|
|
}
|
|
matrix.push(row);
|
|
}
|
|
return matrix;
|
|
}
|
|
|
|
/**
|
|
* Zero-initialized matrix
|
|
* @private
|
|
*/
|
|
_zeros(rows, cols) {
|
|
return Array(rows).fill(null).map(() => Array(cols).fill(0));
|
|
}
|
|
|
|
/**
|
|
* Gaussian random number (Box-Muller transform)
|
|
* @private
|
|
*/
|
|
_gaussianRandom() {
|
|
const u1 = Math.random();
|
|
const u2 = Math.random();
|
|
return Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
|
|
}
|
|
|
|
// ============================================
|
|
// TRAINING METHODS
|
|
// ============================================
|
|
|
|
/**
|
|
* Train the adapter on provided examples
|
|
*
|
|
* @param {TrainingExample[]} examples - Training examples
|
|
* @param {TrainingOptions} [options={}] - Training options
|
|
* @returns {Promise<Object>} Training result with metrics
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* const result = await lora.train([
|
|
* { input: 'translate to python', output: 'def main():' },
|
|
* { input: 'write a loop', output: 'for i in range(10):' }
|
|
* ], {
|
|
* epochs: 10,
|
|
* lr: 1e-4,
|
|
* batchSize: 4,
|
|
* onProgress: (progress) => console.log(`${progress.epoch}/${progress.totalEpochs}`)
|
|
* });
|
|
* console.log(`Final loss: ${result.finalLoss}`);
|
|
* ```
|
|
*/
|
|
async train(examples, options = {}) {
|
|
if (this.isTraining) {
|
|
throw new Error('Training already in progress');
|
|
}
|
|
|
|
const opts = {
|
|
epochs: 10,
|
|
lr: 1e-4,
|
|
batchSize: 4,
|
|
scheduler: 'cosine',
|
|
warmupSteps: 10,
|
|
useEWC: false,
|
|
ewcLambda: 1000,
|
|
gradientCheckpointing: true,
|
|
onProgress: null,
|
|
...options,
|
|
};
|
|
|
|
this.isTraining = true;
|
|
this.emit('training:start', { examples: examples.length, options: opts });
|
|
|
|
const startTime = Date.now();
|
|
const lossHistory = [];
|
|
const totalSteps = Math.ceil(examples.length / opts.batchSize) * opts.epochs;
|
|
let currentStep = 0;
|
|
|
|
try {
|
|
// Preprocess examples
|
|
const processedExamples = await this._preprocessExamples(examples);
|
|
|
|
// Initialize EWC if enabled and we have prior state
|
|
if (opts.useEWC && this.ewcState) {
|
|
this.emit('training:ewc', { lambda: opts.ewcLambda });
|
|
}
|
|
|
|
// Training loop
|
|
for (let epoch = 0; epoch < opts.epochs; epoch++) {
|
|
const epochLosses = [];
|
|
|
|
// Shuffle examples
|
|
const shuffled = [...processedExamples].sort(() => Math.random() - 0.5);
|
|
|
|
// Process batches
|
|
for (let i = 0; i < shuffled.length; i += opts.batchSize) {
|
|
const batch = shuffled.slice(i, i + opts.batchSize);
|
|
currentStep++;
|
|
|
|
// Compute learning rate with scheduler
|
|
const lr = this._computeLearningRate(opts, currentStep, totalSteps);
|
|
|
|
// Forward and backward pass
|
|
const batchLoss = await this._trainBatch(batch, lr, opts);
|
|
epochLosses.push(batchLoss);
|
|
|
|
// EWC regularization
|
|
if (opts.useEWC && this.ewcState) {
|
|
this._applyEWCPenalty(opts.ewcLambda, lr);
|
|
}
|
|
|
|
// Progress callback
|
|
if (opts.onProgress) {
|
|
opts.onProgress({
|
|
epoch,
|
|
totalEpochs: opts.epochs,
|
|
step: currentStep,
|
|
totalSteps,
|
|
loss: batchLoss,
|
|
lr,
|
|
});
|
|
}
|
|
|
|
this.emit('training:step', {
|
|
epoch,
|
|
step: currentStep,
|
|
loss: batchLoss,
|
|
lr,
|
|
});
|
|
}
|
|
|
|
const epochLoss = epochLosses.reduce((a, b) => a + b, 0) / epochLosses.length;
|
|
lossHistory.push(epochLoss);
|
|
|
|
this.emit('training:epoch', {
|
|
epoch,
|
|
loss: epochLoss,
|
|
lossHistory,
|
|
});
|
|
}
|
|
|
|
// Update EWC state for continual learning
|
|
if (opts.useEWC) {
|
|
this.ewcState = this._computeEWCState(processedExamples);
|
|
}
|
|
|
|
// Update stats
|
|
this.stats.totalTrainingSamples += examples.length;
|
|
this.stats.totalTrainingTime += Date.now() - startTime;
|
|
this.metadata.trainingSamples = this.stats.totalTrainingSamples;
|
|
this.metadata.trainingEpochs += opts.epochs;
|
|
|
|
const result = {
|
|
finalLoss: lossHistory[lossHistory.length - 1],
|
|
lossHistory,
|
|
trainingTime: Date.now() - startTime,
|
|
totalSteps: currentStep,
|
|
examples: examples.length,
|
|
epochs: opts.epochs,
|
|
};
|
|
|
|
this.emit('training:complete', result);
|
|
return result;
|
|
|
|
} finally {
|
|
this.isTraining = false;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Train on a single correction (online learning)
|
|
*
|
|
* @param {string} input - Original input
|
|
* @param {string} wrongOutput - Incorrect model output
|
|
* @param {string} correctOutput - User-provided correction
|
|
* @returns {Promise<Object>} Training result
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* // User corrects a model mistake
|
|
* await lora.trainOnCorrection(
|
|
* 'write hello world',
|
|
* 'print(hello world)', // Wrong output
|
|
* 'print("hello world")' // Correct output
|
|
* );
|
|
* ```
|
|
*/
|
|
async trainOnCorrection(input, wrongOutput, correctOutput) {
|
|
this.stats.corrections++;
|
|
|
|
const result = await this.train([
|
|
{ input, output: correctOutput, quality: 1.5 }, // Upweight corrections
|
|
], {
|
|
epochs: 3,
|
|
lr: 5e-4, // Higher LR for quick adaptation
|
|
useEWC: true, // Preserve prior knowledge
|
|
});
|
|
|
|
this.emit('correction:applied', {
|
|
input,
|
|
wrongOutput,
|
|
correctOutput,
|
|
result,
|
|
});
|
|
|
|
return result;
|
|
}
|
|
|
|
/**
|
|
* Preprocess training examples into embeddings
|
|
* @private
|
|
*/
|
|
async _preprocessExamples(examples) {
|
|
const processed = [];
|
|
|
|
for (const example of examples) {
|
|
// Generate input embedding (simple hash-based for now)
|
|
const inputEmb = this._hashEmbed(example.input);
|
|
const outputEmb = this._hashEmbed(example.output);
|
|
|
|
processed.push({
|
|
input: example.input,
|
|
output: example.output,
|
|
inputEmb,
|
|
outputEmb,
|
|
quality: example.quality || 1.0,
|
|
});
|
|
}
|
|
|
|
return processed;
|
|
}
|
|
|
|
/**
|
|
* Hash-based embedding (fallback when no model loaded)
|
|
* @private
|
|
*/
|
|
_hashEmbed(text) {
|
|
const dim = this.config.embeddingDim;
|
|
const hash = createHash('sha256').update(text).digest();
|
|
const embedding = new Float32Array(dim);
|
|
|
|
// Deterministic pseudo-random from hash
|
|
for (let i = 0; i < dim; i++) {
|
|
embedding[i] = (hash[i % 32] - 128) / 128;
|
|
}
|
|
|
|
// Add character-level features
|
|
for (let i = 0; i < text.length && i < dim; i++) {
|
|
embedding[i] += (text.charCodeAt(i) - 64) / 256;
|
|
}
|
|
|
|
// Normalize
|
|
let norm = 0;
|
|
for (let i = 0; i < dim; i++) {
|
|
norm += embedding[i] * embedding[i];
|
|
}
|
|
norm = Math.sqrt(norm) || 1;
|
|
for (let i = 0; i < dim; i++) {
|
|
embedding[i] /= norm;
|
|
}
|
|
|
|
return Array.from(embedding);
|
|
}
|
|
|
|
/**
|
|
* Compute learning rate with scheduler
|
|
* @private
|
|
*/
|
|
_computeLearningRate(opts, step, totalSteps) {
|
|
const baseLR = opts.lr;
|
|
|
|
// Warmup phase
|
|
if (step < opts.warmupSteps) {
|
|
return baseLR * (step / opts.warmupSteps);
|
|
}
|
|
|
|
const decayStep = step - opts.warmupSteps;
|
|
const decayTotal = totalSteps - opts.warmupSteps;
|
|
|
|
switch (opts.scheduler) {
|
|
case 'constant':
|
|
return baseLR;
|
|
case 'linear':
|
|
return baseLR * (1 - decayStep / decayTotal);
|
|
case 'cosine':
|
|
return baseLR * 0.5 * (1 + Math.cos(Math.PI * decayStep / decayTotal));
|
|
case 'exponential':
|
|
return baseLR * Math.pow(0.9, decayStep / 100);
|
|
default:
|
|
return baseLR;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Train on a single batch
|
|
* @private
|
|
*/
|
|
async _trainBatch(batch, lr, opts) {
|
|
let totalLoss = 0;
|
|
|
|
for (const example of batch) {
|
|
// Forward pass through adapters
|
|
const adapted = this._forwardAdapters(example.inputEmb);
|
|
|
|
// Compute loss (MSE between adapted output and target embedding)
|
|
let loss = 0;
|
|
for (let i = 0; i < adapted.length && i < example.outputEmb.length; i++) {
|
|
const diff = adapted[i] - example.outputEmb[i];
|
|
loss += diff * diff;
|
|
}
|
|
loss = (loss / adapted.length) * example.quality;
|
|
|
|
// Backward pass (gradient descent on adapter weights)
|
|
this._backwardAdapters(example.inputEmb, example.outputEmb, lr, opts);
|
|
|
|
totalLoss += loss;
|
|
}
|
|
|
|
return totalLoss / batch.length;
|
|
}
|
|
|
|
/**
|
|
* Forward pass through all adapters
|
|
* @private
|
|
*/
|
|
_forwardAdapters(input) {
|
|
let output = [...input];
|
|
|
|
for (const [, adapter] of this.adapters) {
|
|
output = this._forwardSingleAdapter(output, adapter);
|
|
}
|
|
|
|
return output;
|
|
}
|
|
|
|
/**
|
|
* Forward through a single adapter
|
|
* @private
|
|
*/
|
|
_forwardSingleAdapter(input, adapter) {
|
|
const rank = this.config.rank;
|
|
const dim = Math.min(input.length, this.config.embeddingDim);
|
|
|
|
// input @ A (dim -> rank)
|
|
const hidden = new Float64Array(rank);
|
|
for (let r = 0; r < rank; r++) {
|
|
let sum = 0;
|
|
for (let d = 0; d < dim; d++) {
|
|
sum += input[d] * adapter.loraA[d][r];
|
|
}
|
|
hidden[r] = sum;
|
|
}
|
|
|
|
// hidden @ B (rank -> dim) + residual
|
|
const output = [...input];
|
|
for (let d = 0; d < dim; d++) {
|
|
let delta = 0;
|
|
for (let r = 0; r < rank; r++) {
|
|
delta += hidden[r] * adapter.loraB[r][d];
|
|
}
|
|
output[d] += adapter.scaling * delta;
|
|
}
|
|
|
|
return output;
|
|
}
|
|
|
|
/**
|
|
* Backward pass through all adapters
|
|
* @private
|
|
*/
|
|
_backwardAdapters(input, target, lr, opts) {
|
|
for (const [, adapter] of this.adapters) {
|
|
this._backwardSingleAdapter(input, target, adapter, lr);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Backward through a single adapter
|
|
* @private
|
|
*/
|
|
_backwardSingleAdapter(input, target, adapter, lr) {
|
|
const rank = this.config.rank;
|
|
const dim = Math.min(input.length, this.config.embeddingDim);
|
|
|
|
// Compute forward pass for gradient computation
|
|
const hidden = new Float64Array(rank);
|
|
for (let r = 0; r < rank; r++) {
|
|
for (let d = 0; d < dim; d++) {
|
|
hidden[r] += input[d] * adapter.loraA[d][r];
|
|
}
|
|
}
|
|
|
|
// Compute adapted output
|
|
const adapted = [...input];
|
|
for (let d = 0; d < dim; d++) {
|
|
for (let r = 0; r < rank; r++) {
|
|
adapted[d] += adapter.scaling * hidden[r] * adapter.loraB[r][d];
|
|
}
|
|
}
|
|
|
|
// Compute output gradient (MSE derivative)
|
|
const gradOutput = adapted.map((val, i) =>
|
|
2 * (val - (target[i] || 0)) / dim
|
|
);
|
|
|
|
// Gradient for B: hidden^T @ gradOutput
|
|
for (let r = 0; r < rank; r++) {
|
|
for (let d = 0; d < dim; d++) {
|
|
const grad = hidden[r] * gradOutput[d] * adapter.scaling;
|
|
adapter.loraB[r][d] -= lr * grad;
|
|
}
|
|
}
|
|
|
|
// Gradient for hidden: gradOutput @ B^T
|
|
const gradHidden = new Float64Array(rank);
|
|
for (let r = 0; r < rank; r++) {
|
|
for (let d = 0; d < dim; d++) {
|
|
gradHidden[r] += gradOutput[d] * adapter.loraB[r][d] * adapter.scaling;
|
|
}
|
|
}
|
|
|
|
// Gradient for A: input^T @ gradHidden
|
|
for (let d = 0; d < dim; d++) {
|
|
for (let r = 0; r < rank; r++) {
|
|
const grad = input[d] * gradHidden[r];
|
|
adapter.loraA[d][r] -= lr * grad;
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Compute EWC state (Fisher information matrix approximation)
|
|
* @private
|
|
*/
|
|
_computeEWCState(examples) {
|
|
const state = {
|
|
fisherDiag: new Map(),
|
|
optimalParams: new Map(),
|
|
};
|
|
|
|
for (const [name, adapter] of this.adapters) {
|
|
// Store optimal parameters
|
|
state.optimalParams.set(name, {
|
|
loraA: adapter.loraA.map(row => [...row]),
|
|
loraB: adapter.loraB.map(row => [...row]),
|
|
});
|
|
|
|
// Approximate Fisher diagonal with gradient squares
|
|
const fisherA = this._zeros(adapter.loraA.length, adapter.loraA[0].length);
|
|
const fisherB = this._zeros(adapter.loraB.length, adapter.loraB[0].length);
|
|
|
|
// Accumulate gradients squared
|
|
for (const example of examples) {
|
|
const grads = this._computeGradients(example.inputEmb, example.outputEmb, adapter);
|
|
for (let i = 0; i < fisherA.length; i++) {
|
|
for (let j = 0; j < fisherA[0].length; j++) {
|
|
fisherA[i][j] += grads.gradA[i][j] * grads.gradA[i][j];
|
|
}
|
|
}
|
|
for (let i = 0; i < fisherB.length; i++) {
|
|
for (let j = 0; j < fisherB[0].length; j++) {
|
|
fisherB[i][j] += grads.gradB[i][j] * grads.gradB[i][j];
|
|
}
|
|
}
|
|
}
|
|
|
|
// Normalize
|
|
const n = examples.length;
|
|
for (let i = 0; i < fisherA.length; i++) {
|
|
for (let j = 0; j < fisherA[0].length; j++) {
|
|
fisherA[i][j] /= n;
|
|
}
|
|
}
|
|
for (let i = 0; i < fisherB.length; i++) {
|
|
for (let j = 0; j < fisherB[0].length; j++) {
|
|
fisherB[i][j] /= n;
|
|
}
|
|
}
|
|
|
|
state.fisherDiag.set(name, { fisherA, fisherB });
|
|
}
|
|
|
|
return state;
|
|
}
|
|
|
|
/**
|
|
* Compute gradients for a single example
|
|
* @private
|
|
*/
|
|
_computeGradients(input, target, adapter) {
|
|
const rank = this.config.rank;
|
|
const dim = Math.min(input.length, this.config.embeddingDim);
|
|
|
|
// Forward pass
|
|
const hidden = new Float64Array(rank);
|
|
for (let r = 0; r < rank; r++) {
|
|
for (let d = 0; d < dim; d++) {
|
|
hidden[r] += input[d] * adapter.loraA[d][r];
|
|
}
|
|
}
|
|
|
|
const adapted = [...input];
|
|
for (let d = 0; d < dim; d++) {
|
|
for (let r = 0; r < rank; r++) {
|
|
adapted[d] += adapter.scaling * hidden[r] * adapter.loraB[r][d];
|
|
}
|
|
}
|
|
|
|
// Output gradient
|
|
const gradOutput = adapted.map((val, i) =>
|
|
2 * (val - (target[i] || 0)) / dim
|
|
);
|
|
|
|
// Gradient for B
|
|
const gradB = this._zeros(rank, dim);
|
|
for (let r = 0; r < rank; r++) {
|
|
for (let d = 0; d < dim; d++) {
|
|
gradB[r][d] = hidden[r] * gradOutput[d] * adapter.scaling;
|
|
}
|
|
}
|
|
|
|
// Gradient for hidden
|
|
const gradHidden = new Float64Array(rank);
|
|
for (let r = 0; r < rank; r++) {
|
|
for (let d = 0; d < dim; d++) {
|
|
gradHidden[r] += gradOutput[d] * adapter.loraB[r][d] * adapter.scaling;
|
|
}
|
|
}
|
|
|
|
// Gradient for A
|
|
const gradA = this._zeros(dim, rank);
|
|
for (let d = 0; d < dim; d++) {
|
|
for (let r = 0; r < rank; r++) {
|
|
gradA[d][r] = input[d] * gradHidden[r];
|
|
}
|
|
}
|
|
|
|
return { gradA, gradB };
|
|
}
|
|
|
|
/**
|
|
* Apply EWC penalty to prevent catastrophic forgetting
|
|
* @private
|
|
*/
|
|
_applyEWCPenalty(lambda, lr) {
|
|
if (!this.ewcState) return;
|
|
|
|
for (const [name, adapter] of this.adapters) {
|
|
const fisher = this.ewcState.fisherDiag.get(name);
|
|
const optimal = this.ewcState.optimalParams.get(name);
|
|
|
|
if (!fisher || !optimal) continue;
|
|
|
|
// Apply EWC penalty: lambda/2 * F * (theta - theta*)^2
|
|
for (let i = 0; i < adapter.loraA.length; i++) {
|
|
for (let j = 0; j < adapter.loraA[0].length; j++) {
|
|
const diff = adapter.loraA[i][j] - optimal.loraA[i][j];
|
|
adapter.loraA[i][j] -= lr * lambda * fisher.fisherA[i][j] * diff;
|
|
}
|
|
}
|
|
for (let i = 0; i < adapter.loraB.length; i++) {
|
|
for (let j = 0; j < adapter.loraB[0].length; j++) {
|
|
const diff = adapter.loraB[i][j] - optimal.loraB[i][j];
|
|
adapter.loraB[i][j] -= lr * lambda * fisher.fisherB[i][j] * diff;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================
|
|
// INFERENCE METHODS
|
|
// ============================================
|
|
|
|
/**
|
|
* Generate text with the adapted model
|
|
*
|
|
* @param {string} prompt - Input prompt
|
|
* @param {GenerationOptions} [options={}] - Generation options
|
|
* @returns {Promise<Object>} Generation result
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* const result = await lora.generate('Write a Python function', {
|
|
* maxTokens: 128,
|
|
* temperature: 0.8
|
|
* });
|
|
* console.log(result.text);
|
|
* ```
|
|
*/
|
|
async generate(prompt, options = {}) {
|
|
const opts = {
|
|
maxTokens: 64,
|
|
temperature: 0.7,
|
|
topP: 0.9,
|
|
topK: 50,
|
|
repetitionPenalty: 1.1,
|
|
...options,
|
|
};
|
|
|
|
this.stats.totalInferences++;
|
|
|
|
// Embed the prompt
|
|
const promptEmb = this._hashEmbed(prompt);
|
|
|
|
// Apply adapters
|
|
const adapted = this._forwardAdapters(promptEmb);
|
|
|
|
// For now, return simulated generation
|
|
// In production, this would interface with actual ONNX inference
|
|
const result = {
|
|
text: this._simulateGeneration(prompt, adapted, opts),
|
|
prompt,
|
|
adapted: true,
|
|
adaptersApplied: this.adapters.size,
|
|
model: this.baseModelKey,
|
|
options: opts,
|
|
timestamp: Date.now(),
|
|
};
|
|
|
|
this.emit('inference:complete', result);
|
|
return result;
|
|
}
|
|
|
|
/**
|
|
* Simulate text generation (placeholder for actual LLM inference)
|
|
* @private
|
|
*/
|
|
_simulateGeneration(prompt, embedding, opts) {
|
|
// This is a placeholder - in production, would use ONNX runtime
|
|
// For now, return a template response based on adapter modifications
|
|
const embMagnitude = Math.sqrt(embedding.reduce((sum, v) => sum + v * v, 0));
|
|
|
|
return `[MicroLoRA Adapted Output]\n` +
|
|
`Prompt: ${prompt.slice(0, 50)}...\n` +
|
|
`Embedding magnitude: ${embMagnitude.toFixed(4)}\n` +
|
|
`Adapters: ${this.adapters.size} active\n` +
|
|
`Model: ${this.baseModelKey}`;
|
|
}
|
|
|
|
/**
|
|
* Embed text using the model
|
|
*
|
|
* @param {string} text - Text to embed
|
|
* @returns {Promise<number[]>} Embedding vector
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* const embedding = await lora.embed('Hello world');
|
|
* console.log(`Dimension: ${embedding.length}`);
|
|
* ```
|
|
*/
|
|
async embed(text) {
|
|
const baseEmb = this._hashEmbed(text);
|
|
const adapted = this._forwardAdapters(baseEmb);
|
|
return adapted;
|
|
}
|
|
|
|
// ============================================
|
|
// ADAPTER MANAGEMENT
|
|
// ============================================
|
|
|
|
/**
|
|
* Save adapter to file or return serialized data
|
|
*
|
|
* @param {string} [path] - Optional file path (Node.js only)
|
|
* @returns {Promise<Object>} Serialized adapter data
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* // Save to file (Node.js)
|
|
* await lora.saveAdapter('./my-adapter.json');
|
|
*
|
|
* // Get serialized data (browser)
|
|
* const data = await lora.saveAdapter();
|
|
* localStorage.setItem('my-adapter', JSON.stringify(data));
|
|
* ```
|
|
*/
|
|
async saveAdapter(path) {
|
|
const data = {
|
|
version: '1.0.0',
|
|
format: 'microlora',
|
|
metadata: { ...this.metadata },
|
|
config: { ...this.config },
|
|
baseModel: this.baseModelKey,
|
|
adapters: {},
|
|
stats: { ...this.stats },
|
|
createdAt: this.metadata.createdAt,
|
|
savedAt: Date.now(),
|
|
};
|
|
|
|
// Serialize adapters
|
|
for (const [name, adapter] of this.adapters) {
|
|
data.adapters[name] = {
|
|
loraA: adapter.loraA,
|
|
loraB: adapter.loraB,
|
|
scaling: adapter.scaling,
|
|
};
|
|
}
|
|
|
|
// Save to file if path provided (Node.js)
|
|
if (path && typeof process !== 'undefined') {
|
|
const fs = await import('fs/promises');
|
|
await fs.writeFile(path, JSON.stringify(data, null, 2));
|
|
this.emit('adapter:saved', { path, size: JSON.stringify(data).length });
|
|
}
|
|
|
|
return data;
|
|
}
|
|
|
|
/**
|
|
* Load adapter from file or serialized data
|
|
*
|
|
* @param {string|Object} pathOrData - File path or serialized adapter data
|
|
* @returns {Promise<void>}
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* // Load from file (Node.js)
|
|
* await lora.loadAdapter('./my-adapter.json');
|
|
*
|
|
* // Load from data (browser)
|
|
* const data = JSON.parse(localStorage.getItem('my-adapter'));
|
|
* await lora.loadAdapter(data);
|
|
* ```
|
|
*/
|
|
async loadAdapter(pathOrData) {
|
|
let data;
|
|
|
|
if (typeof pathOrData === 'string') {
|
|
// Load from file (Node.js)
|
|
const fs = await import('fs/promises');
|
|
const content = await fs.readFile(pathOrData, 'utf-8');
|
|
data = JSON.parse(content);
|
|
} else {
|
|
data = pathOrData;
|
|
}
|
|
|
|
// Validate format
|
|
if (data.format !== 'microlora') {
|
|
throw new Error(`Unsupported adapter format: ${data.format}`);
|
|
}
|
|
|
|
// Check base model compatibility
|
|
if (data.baseModel !== this.baseModelKey) {
|
|
console.warn(`Warning: Adapter was trained on ${data.baseModel}, ` +
|
|
`but loading into ${this.baseModelKey}`);
|
|
}
|
|
|
|
// Load adapter weights
|
|
this.adapters.clear();
|
|
for (const [name, adapter] of Object.entries(data.adapters)) {
|
|
this.adapters.set(name, {
|
|
loraA: adapter.loraA,
|
|
loraB: adapter.loraB,
|
|
scaling: adapter.scaling,
|
|
});
|
|
}
|
|
|
|
// Restore metadata
|
|
this.metadata = { ...this.metadata, ...data.metadata };
|
|
this.stats = { ...this.stats, ...data.stats };
|
|
|
|
this.emit('adapter:loaded', {
|
|
path: typeof pathOrData === 'string' ? pathOrData : null,
|
|
adapters: this.adapters.size,
|
|
metadata: this.metadata,
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Merge multiple adapters with weights
|
|
*
|
|
* @param {Array<{adapter: Object, weight: number}>} adapters - Adapters with weights
|
|
* @returns {Promise<void>}
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* await lora.mergeAdapters([
|
|
* { adapter: codeAdapter, weight: 0.7 },
|
|
* { adapter: mathAdapter, weight: 0.3 }
|
|
* ]);
|
|
* ```
|
|
*/
|
|
async mergeAdapters(adapters) {
|
|
if (adapters.length === 0) return;
|
|
|
|
// Normalize weights
|
|
const totalWeight = adapters.reduce((sum, a) => sum + a.weight, 0);
|
|
const normalizedAdapters = adapters.map(a => ({
|
|
...a,
|
|
weight: a.weight / totalWeight,
|
|
}));
|
|
|
|
// Merge each module
|
|
for (const [name, currentAdapter] of this.adapters) {
|
|
// Reset to zero
|
|
for (let i = 0; i < currentAdapter.loraA.length; i++) {
|
|
for (let j = 0; j < currentAdapter.loraA[0].length; j++) {
|
|
currentAdapter.loraA[i][j] = 0;
|
|
}
|
|
}
|
|
for (let i = 0; i < currentAdapter.loraB.length; i++) {
|
|
for (let j = 0; j < currentAdapter.loraB[0].length; j++) {
|
|
currentAdapter.loraB[i][j] = 0;
|
|
}
|
|
}
|
|
|
|
// Weighted sum of adapters
|
|
for (const { adapter, weight } of normalizedAdapters) {
|
|
const adapterWeights = adapter.adapters?.[name] || adapter.adapters?.get?.(name);
|
|
if (!adapterWeights) continue;
|
|
|
|
for (let i = 0; i < currentAdapter.loraA.length; i++) {
|
|
for (let j = 0; j < currentAdapter.loraA[0].length; j++) {
|
|
currentAdapter.loraA[i][j] += weight * (adapterWeights.loraA[i]?.[j] || 0);
|
|
}
|
|
}
|
|
for (let i = 0; i < currentAdapter.loraB.length; i++) {
|
|
for (let j = 0; j < currentAdapter.loraB[0].length; j++) {
|
|
currentAdapter.loraB[i][j] += weight * (adapterWeights.loraB[i]?.[j] || 0);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
this.stats.adaptations++;
|
|
this.emit('adapter:merged', {
|
|
count: adapters.length,
|
|
weights: normalizedAdapters.map(a => a.weight),
|
|
});
|
|
}
|
|
|
|
// ============================================
|
|
// EXPORT METHODS
|
|
// ============================================
|
|
|
|
/**
|
|
* Export adapter to ONNX format
|
|
*
|
|
* @returns {Promise<Uint8Array>} ONNX model bytes
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* const onnxBytes = await lora.exportToONNX();
|
|
* // Save or deploy the ONNX model
|
|
* ```
|
|
*/
|
|
async exportToONNX() {
|
|
// Build ONNX-compatible structure
|
|
const onnxData = {
|
|
format: 'onnx-lora-adapter',
|
|
version: '1.0',
|
|
baseModel: this.baseModelKey,
|
|
config: this.config,
|
|
adapters: [],
|
|
};
|
|
|
|
for (const [name, adapter] of this.adapters) {
|
|
onnxData.adapters.push({
|
|
name,
|
|
loraA: {
|
|
shape: [adapter.loraA.length, adapter.loraA[0].length],
|
|
data: adapter.loraA.flat(),
|
|
},
|
|
loraB: {
|
|
shape: [adapter.loraB.length, adapter.loraB[0].length],
|
|
data: adapter.loraB.flat(),
|
|
},
|
|
scaling: adapter.scaling,
|
|
});
|
|
}
|
|
|
|
// Convert to protobuf-like format (simplified)
|
|
const json = JSON.stringify(onnxData);
|
|
const bytes = new TextEncoder().encode(json);
|
|
|
|
this.emit('export:onnx', { size: bytes.length });
|
|
return bytes;
|
|
}
|
|
|
|
/**
|
|
* Export adapter to IPFS-compatible format
|
|
*
|
|
* @returns {Promise<Object>} IPFS-ready data with CID placeholder
|
|
*
|
|
* @example
|
|
* ```javascript
|
|
* const ipfsData = await lora.exportToIPFS();
|
|
* // Upload to IPFS using your preferred client
|
|
* const cid = await ipfsClient.add(ipfsData.content);
|
|
* ```
|
|
*/
|
|
async exportToIPFS() {
|
|
const adapterData = await this.saveAdapter();
|
|
|
|
// Add IPFS-specific metadata
|
|
const ipfsData = {
|
|
...adapterData,
|
|
ipfs: {
|
|
version: 1,
|
|
contentType: 'application/json',
|
|
compression: 'none',
|
|
chunks: 1,
|
|
},
|
|
};
|
|
|
|
const content = JSON.stringify(ipfsData);
|
|
const hash = createHash('sha256').update(content).digest('hex');
|
|
|
|
return {
|
|
content: new TextEncoder().encode(content),
|
|
hash,
|
|
size: content.length,
|
|
metadata: {
|
|
name: this.metadata.name,
|
|
description: this.metadata.description,
|
|
domain: this.metadata.domain,
|
|
baseModel: this.baseModelKey,
|
|
},
|
|
};
|
|
}
|
|
|
|
// ============================================
|
|
// UTILITY METHODS
|
|
// ============================================
|
|
|
|
/**
|
|
* Get adapter metadata
|
|
*
|
|
* @returns {AdapterMetadata} Current adapter metadata
|
|
*/
|
|
getMetadata() {
|
|
return { ...this.metadata };
|
|
}
|
|
|
|
/**
|
|
* Set adapter metadata
|
|
*
|
|
* @param {Partial<AdapterMetadata>} metadata - Metadata to update
|
|
*/
|
|
setMetadata(metadata) {
|
|
this.metadata = { ...this.metadata, ...metadata };
|
|
}
|
|
|
|
/**
|
|
* Get training and inference statistics
|
|
*
|
|
* @returns {Object} Current statistics
|
|
*/
|
|
getStats() {
|
|
return {
|
|
...this.stats,
|
|
adapters: this.adapters.size,
|
|
config: { ...this.config },
|
|
baseModel: this.baseModelKey,
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Reset adapter to initial state
|
|
*/
|
|
reset() {
|
|
this.adapters = this._initializeAdapters();
|
|
this.ewcState = null;
|
|
this.trainingState = null;
|
|
this.stats = {
|
|
totalTrainingSamples: 0,
|
|
totalTrainingTime: 0,
|
|
totalInferences: 0,
|
|
corrections: 0,
|
|
adaptations: 0,
|
|
};
|
|
this.emit('adapter:reset');
|
|
}
|
|
|
|
/**
|
|
* Get number of trainable parameters
|
|
*
|
|
* @returns {number} Total trainable parameters
|
|
*/
|
|
numParameters() {
|
|
let total = 0;
|
|
for (const [, adapter] of this.adapters) {
|
|
total += adapter.loraA.length * adapter.loraA[0].length;
|
|
total += adapter.loraB.length * adapter.loraB[0].length;
|
|
}
|
|
return total;
|
|
}
|
|
|
|
/**
|
|
* Check if adapter has been trained
|
|
*
|
|
* @returns {boolean} True if any training has occurred
|
|
*/
|
|
isTrained() {
|
|
return this.stats.totalTrainingSamples > 0;
|
|
}
|
|
}
|
|
|
|
// ============================================
|
|
// EXPORTS
|
|
// ============================================
|
|
|
|
export default MicroLoRA;
|