feat: add text embedding model and pipeline.

Author: daviddaytw

src/__tests__/text-embedding.model.test.tsx

@@ -0,0 +1,58 @@
+import { TextEmbedding } from "../models/text-embedding";
+import { InferenceSession } from "onnxruntime-react-native";
+
+describe("TextEmbedding Model", () => {
+  let model: TextEmbedding;
+
+  beforeEach(() => {
+    model = new TextEmbedding();
+  });
+
+  afterEach(async () => {
+    await model.release();
+  });
+
+  it("should initialize properly", () => {
+    expect(model).toBeInstanceOf(TextEmbedding);
+  });
+
+  it("should throw error when session is undefined", async () => {
+    await expect(model.embed([1n, 2n, 3n])).rejects.toThrow(
+      "Session is undefined",
+    );
+  });
+
+  it("should throw error when no embedding output is found", async () => {
+    // Mock session run to return empty outputs
+    const mockRun = jest.fn().mockResolvedValue({});
+    (model as any).sess = {
+      run: mockRun,
+      release: jest.fn().mockResolvedValue(undefined),
+    } as Partial<InferenceSession>;
+
+    await expect(model.embed([1n, 2n, 3n])).rejects.toThrow(
+      "No embedding output found in model outputs",
+    );
+  });
+
+  it("should properly calculate mean embeddings", async () => {
+    // Mock session run to return sample embeddings
+    const mockEmbeddings = new Float32Array([1, 2, 3, 4, 5, 6]); // 2 tokens, 3 dimensions
+    const mockRun = jest.fn().mockResolvedValue({
+      last_hidden_state: {
+        data: mockEmbeddings,
+        dims: [1, 2, 3], // [batch_size, sequence_length, hidden_size]
+      },
+    });
+    (model as any).sess = {
+      run: mockRun,
+      release: jest.fn().mockResolvedValue(undefined),
+    } as Partial<InferenceSession>;
+
+    const result = await model.embed([1n, 2n]);
+
+    // Expected mean values: [2.5, 3.5, 4.5]
+    expect(Array.from(result)).toEqual([2.5, 3.5, 4.5]);
+    expect(mockRun).toHaveBeenCalled();
+  });
+});

src/__tests__/text-embedding.pipeline.test.tsx

@@ -0,0 +1,59 @@
+import TextEmbeddingPipeline from "../pipelines/text-embedding";
+
+// Mock the TextEmbedding model
+jest.mock("../models/text-embedding", () => {
+  return {
+    TextEmbedding: jest.fn().mockImplementation(() => ({
+      load: jest.fn().mockResolvedValue(undefined),
+      embed: jest.fn().mockResolvedValue(new Float32Array([0.1, 0.2, 0.3])),
+      release: jest.fn().mockResolvedValue(undefined),
+    })),
+  };
+});
+
+// Create a callable tokenizer mock
+const createCallableTokenizer = () => {
+  const tokenizer = jest.fn().mockResolvedValue({
+    input_ids: [1n, 2n, 3n],
+  });
+  return tokenizer;
+};
+
+jest.mock("@xenova/transformers", () => ({
+  env: {
+    allowRemoteModels: true,
+    allowLocalModels: false,
+  },
+  AutoTokenizer: {
+    from_pretrained: jest.fn().mockResolvedValue(createCallableTokenizer()),
+  },
+}));
+
+describe("TextEmbedding Pipeline", () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  afterEach(async () => {
+    await TextEmbeddingPipeline.release();
+  });
+
+  it("should throw error when not initialized", async () => {
+    await expect(TextEmbeddingPipeline.embed("test text")).rejects.toThrow(
+      "Tokenizer undefined, please initialize first",
+    );
+  });
+
+  it("should initialize properly", async () => {
+    await expect(
+      TextEmbeddingPipeline.init("test-model", "model.onnx"),
+    ).resolves.not.toThrow();
+  });
+
+  it("should generate embeddings", async () => {
+    await TextEmbeddingPipeline.init("test-model", "model.onnx");
+    const embeddings = await TextEmbeddingPipeline.embed("test text");
+    expect(embeddings).toBeInstanceOf(Float32Array);
+    expect(embeddings.length).toBe(3);
+  });
+});

src/index.tsx

@@ -1,12 +1,16 @@
 import { TextGeneration } from "./models/text-generation";
+import { TextEmbedding } from "./models/text-embedding";
 import TextGenerationPipeline from "./pipelines/text-generation";
+import TextEmbeddingPipeline from "./pipelines/text-embedding";
 
 export const Pipeline = {
   TextGeneration: TextGenerationPipeline,
+  TextEmbedding: TextEmbeddingPipeline,
 };
 
 export const Model = {
   TextGeneration,
+  TextEmbedding,
 };
 
 export default {
@@ -16,4 +20,6 @@ export default {
 
 export type * from "./models/base";
 export type * from "./models/text-generation";
+export type * from "./models/text-embedding";
 export type * from "./pipelines/text-generation";
+export type * from "./pipelines/text-embedding";

src/models/text-embedding.tsx

@@ -0,0 +1,61 @@
+import "text-encoding-polyfill";
+import { Tensor } from "onnxruntime-react-native";
+import { Base } from "./base";
+
+/**
+ * Class to handle text embedding model on top of onnxruntime
+ */
+export class TextEmbedding extends Base {
+  /**
+   * Generate embeddings from input tokens
+   *
+   * @param tokens Input tokens to generate embeddings from
+   * @returns Float32Array containing the embedding vector
+   */
+  public async embed(tokens: bigint[]): Promise<Float32Array> {
+    const feed = this.feed;
+    const inputIdsTensor = new Tensor(
+      "int64",
+      BigInt64Array.from(tokens.map(BigInt)),
+      [1, tokens.length],
+    );
+    feed.input_ids = inputIdsTensor;
+
+    // Create attention mask (1 for all tokens)
+    feed.attention_mask = new Tensor(
+      "int64",
+      BigInt64Array.from({ length: tokens.length }, () => 1n),
+      [1, tokens.length],
+    );
+
+    if (!this.sess) {
+      throw new Error("Session is undefined");
+    }
+
+    // Run inference to get embeddings
+    const outputs = await this.sess.run(feed);
+
+    // The model typically outputs the embeddings as 'last_hidden_state' or 'embeddings'
+    // We take the mean of the token embeddings to get a single vector
+    const embeddings = outputs.last_hidden_state || outputs.embeddings;
+
+    if (!embeddings) {
+      throw new Error("No embedding output found in model outputs");
+    }
+
+    // Calculate mean across token dimension (dim 1) to get a single embedding vector
+    const data = embeddings.data as Float32Array;
+    const [, seqLen, hiddenSize] = embeddings.dims;
+    const result = new Float32Array(hiddenSize);
+
+    for (let h = 0; h < hiddenSize; h++) {
+      let sum = 0;
+      for (let s = 0; s < seqLen; s++) {
+        sum += data[s * hiddenSize + h];
+      }
+      result[h] = sum / seqLen;
+    }
+
+    return result;
+  }
+}

src/pipelines/text-embedding.tsx

@@ -0,0 +1,79 @@
+import { env, AutoTokenizer, PreTrainedTokenizer } from "@xenova/transformers";
+import { TextEmbedding as Model } from "../models/text-embedding";
+import { LoadOptions } from "../models/base";
+
+/** Initialization Options for Text Embedding */
+export interface TextEmbeddingOptions extends LoadOptions {
+  /** Shows special tokens in the output. */
+  show_special: boolean;
+}
+
+// Set up environment for transformers.js tokenizer
+env.allowRemoteModels = true;
+env.allowLocalModels = false;
+
+// Declare tokenizer and model
+let tokenizer: PreTrainedTokenizer;
+const model: Model = new Model();
+
+// Initialize options with default values
+let _options: TextEmbeddingOptions = {
+  show_special: false,
+  max_tokens: 512, // typical max length for embedding models
+  fetch: async (url) => url,
+  verbose: false,
+  externalData: false,
+  executionProviders: ["cpu"],
+};
+
+/**
+ * Generates embeddings from the input text.
+ *
+ * @param text - The input text to generate embeddings from.
+ * @returns Float32Array containing the embedding vector.
+ */
+async function embed(text: string): Promise<Float32Array> {
+  if (!tokenizer) {
+    throw new Error("Tokenizer undefined, please initialize first.");
+  }
+
+  const { input_ids } = await tokenizer(text, {
+    return_tensor: false,
+    padding: true,
+    truncation: true,
+    max_length: _options.max_tokens,
+  });
+
+  return await model.embed(input_ids);
+}
+
+/**
+ * Loads the model and tokenizer with the specified options.
+ *
+ * @param model_name - The name of the model to load.
+ * @param onnx_path - The path to the ONNX model.
+ * @param options - Optional initialization options.
+ */
+async function init(
+  model_name: string,
+  onnx_path: string,
+  options?: Partial<TextEmbeddingOptions>,
+): Promise<void> {
+  _options = { ..._options, ...options };
+  tokenizer = await AutoTokenizer.from_pretrained(model_name);
+  await model.load(model_name, onnx_path, _options);
+}
+
+/**
+ * Releases the resources used by the model.
+ */
+async function release(): Promise<void> {
+  await model.release();
+}
+
+// Export functions for external use
+export default {
+  init,
+  embed,
+  release,
+};