Spaces:
Running
Running
File size: 1,900 Bytes
519a20c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
export class WebLlmVectorProvider {
/** @type {object?} WebLLM engine */
#engine = null;
constructor() {
this.#engine = null;
}
/**
* Check if WebLLM is available and up-to-date
* @throws {Error} If WebLLM is not available or not up-to-date
*/
#checkWebLlm() {
if (!Object.hasOwn(SillyTavern, 'llm')) {
throw new Error('WebLLM is not available', { cause: 'webllm-not-available' });
}
if (typeof SillyTavern.llm.generateEmbedding !== 'function') {
throw new Error('WebLLM is not updated', { cause: 'webllm-not-updated' });
}
}
/**
* Initialize the engine with a model.
* @param {string} modelId Model ID to initialize the engine with
* @returns {Promise<void>} Promise that resolves when the engine is initialized
*/
#initEngine(modelId) {
this.#checkWebLlm();
if (!this.#engine) {
this.#engine = SillyTavern.llm.getEngine();
}
return this.#engine.loadModel(modelId);
}
/**
* Get available models.
* @returns {{id:string, toString: function(): string}[]} Array of available models
*/
getModels() {
this.#checkWebLlm();
return SillyTavern.llm.getEmbeddingModels();
}
/**
* Generate embeddings for a list of texts.
* @param {string[]} texts Array of texts to generate embeddings for
* @param {string} modelId Model to use for generating embeddings
* @returns {Promise<number[][]>} Array of embeddings for each text
*/
async embedTexts(texts, modelId) {
await this.#initEngine(modelId);
return this.#engine.generateEmbedding(texts);
}
/**
* Loads a model into the engine.
* @param {string} modelId Model ID to load
*/
async loadModel(modelId) {
await this.#initEngine(modelId);
}
}
|