askexperts/dist/rag/XenovaEmbeddings.js

AskExperts SDK: build and use AI experts - ask them questions and pay with bitcoin on an open protocol

// @ts-ignore
import { pipeline } from "@xenova/transformers";
/**
 * Implementation of RagEmbeddings using Xenova transformers library.
 */
export class XenovaEmbeddings {
    /**
     * Creates a new instance of XenovaEmbeddings.
     * @param model The model name to use for embeddings
     * @param chunkSize The size of text chunks to create
     * @param chunkOverlap The amount of overlap between chunks
     */
    constructor(
    // Good multi-lingual model w/ 256-token context size
    model = XenovaEmbeddings.DEFAULT_MODEL, // 'nomic-ai/nomic-embed-text-v1'
    // I guess we should split by sentence?
    chunkSize = 400, chunkOverlap = 50) {
        this.model = model;
        this.chunkSize = chunkSize;
        this.chunkOverlap = chunkOverlap;
    }
    /**
     * Initializes the embedder pipeline.
     * Must be called before using the embed method.
     * @throws Error if already initialized
     */
    async start() {
        if (this.embedder) {
            throw new Error("Embedder is already initialized");
        }
        this.embedder = await pipeline("feature-extraction", this.model);
    }
    /**
     * Returns the name of the model used for embeddings
     * @returns The model name
     */
    getModelName() {
        return this.model;
    }
    async getVectorSize() {
        return (await this.embedText('')).length;
    }
    /**
     * Splits text into chunks with specified size and overlap.
     * @param text The text to split
     * @returns Array of chunks
     */
    splitTextIntoChunks(text) {
        const chunks = [];
        let index = 0;
        for (let i = 0; i < text.length; i += this.chunkSize - this.chunkOverlap) {
            const end = Math.min(i + this.chunkSize, text.length);
            chunks.push({
                index,
                offset: i,
                text: text.substring(i, end),
                embedding: [],
            });
            index++;
            // If we've reached the end of the text, break
            if (end === text.length)
                break;
        }
        return chunks;
    }
    /**
     * Embeds a single chunk of text.
     * @param text The text to embed
     * @returns The embedding vector
     * @throws Error if start() has not been called
     */
    async embedText(text) {
        if (!this.embedder) {
            throw new Error("Embedder is not initialized. Call start() first.");
        }
        try {
            const embeddings = await this.embedder(text, {
                pooling: "mean", // average over token embeddings
                normalize: true,
            });
            // Convert to regular array if it's not already
            return Array.from(embeddings.data);
        }
        catch (e) {
            console.error(e);
            throw e;
        }
    }
    /**
     * Embeds the given text by splitting it into chunks and generating embeddings.
     * @param text The text to embed
     * @returns Promise resolving to an array of chunks with their embeddings
     */
    async embed(text) {
        const chunks = this.splitTextIntoChunks(text);
        // Process chunks in parallel for better performance
        const chunksWithEmbeddings = await Promise.all(chunks.map(async (chunk) => {
            const embedding = await this.embedText(chunk.text);
            // Create a new chunk with the embedding property
            const chunkWithEmbedding = {
                ...chunk,
                embedding,
            };
            return chunkWithEmbedding;
        }));
        return chunksWithEmbeddings;
    }
}
XenovaEmbeddings.DEFAULT_MODEL = "Xenova/all-MiniLM-L6-v2";
//# sourceMappingURL=XenovaEmbeddings.js.map