From a2a55e21dc0b756fecefd2fde7bb9774a9c60e2a Mon Sep 17 00:00:00 2001 From: Oguz Vuruskaner Date: Sat, 1 Jun 2024 01:53:46 +0300 Subject: [PATCH] community[minor]: DeepInfra embeddings integration #1 (#5382) * Init * fix(type errors) * feat(deepinfra embeddings) * fix(default model) * fix(deepinfra): axios is removed * ref(deepinfra): remove redundant cast * format(deepinfra) * doc(deepinfra) * doc(deepinfra) * Update deepinfra.mdx * Format --------- Co-authored-by: Jacob Lee --- .../integrations/text_embedding/deepinfra.mdx | 127 ++++++++++++ examples/src/embeddings/deepinfra.ts | 12 ++ examples/src/models/embeddings/deepinfra.ts | 12 ++ libs/langchain-community/.gitignore | 4 + libs/langchain-community/langchain.config.js | 1 + libs/langchain-community/package.json | 13 ++ .../src/embeddings/deepinfra.ts | 181 ++++++++++++++++++ .../embeddings/tests/deepinfra.int.test.ts | 34 ++++ 8 files changed, 384 insertions(+) create mode 100644 docs/core_docs/docs/integrations/text_embedding/deepinfra.mdx create mode 100644 examples/src/embeddings/deepinfra.ts create mode 100644 examples/src/models/embeddings/deepinfra.ts create mode 100644 libs/langchain-community/src/embeddings/deepinfra.ts create mode 100644 libs/langchain-community/src/embeddings/tests/deepinfra.int.test.ts diff --git a/docs/core_docs/docs/integrations/text_embedding/deepinfra.mdx b/docs/core_docs/docs/integrations/text_embedding/deepinfra.mdx new file mode 100644 index 000000000000..34b6b942bd7a --- /dev/null +++ b/docs/core_docs/docs/integrations/text_embedding/deepinfra.mdx @@ -0,0 +1,127 @@ +--- +sidebar_label: DeepInfra +--- + +# DeepInfra Embeddings + +The `DeepInfraEmbeddings` class utilizes the DeepInfra API to generate embeddings for given text inputs. This guide will walk you through the setup and usage of the `DeepInfraEmbeddings` class, helping you integrate it into your project seamlessly. + +## Installation + +Install the `@langchain/community` package as shown below: + +import IntegrationInstallTooltip from "@mdx_components/integration_install_tooltip.mdx"; + + + +```bash npm2yarn +npm i @langchain/community +``` + +## Initialization + +With this integration, you can use the DeepInfra embeddings model to get embeddings for your text data. Here is the [link](https://deepinfra.com/models/embeddings) to the embeddings models. + +First, you need to sign up on the DeepInfra website and get the API token from [here](https://deepinfra.com/dash/api_keys). You can copy names from the model cards and start using them in your code. + +To use the `DeepInfraEmbeddings` class, you need an API token from DeepInfra. You can pass this token directly to the constructor or set it as an environment variable (`DEEPINFRA_API_TOKEN`). + +### Basic Usage + +Here’s how to create an instance of `DeepInfraEmbeddings`: + +```typescript +import { DeepInfraEmbeddings } from "@langchain/community/embeddings/deepinfra"; + +const embeddings = new DeepInfraEmbeddings({ + apiToken: "YOUR_API_TOKEN", + modelName: "sentence-transformers/clip-ViT-B-32", // Optional, defaults to "sentence-transformers/clip-ViT-B-32" + batchSize: 1024, // Optional, defaults to 1024 +}); +``` + +If the `apiToken` is not provided, it will be read from the `DEEPINFRA_API_TOKEN` environment variable. + +## Generating Embeddings + +### Embedding a Single Query + +To generate embeddings for a single text query, use the `embedQuery` method: + +```typescript +const embedding = await embeddings.embedQuery( + "What would be a good company name for a company that makes colorful socks?" +); +console.log(embedding); +``` + +### Embedding Multiple Documents + +To generate embeddings for multiple documents, use the `embedDocuments` method. This method will handle batching automatically based on the `batchSize` parameter: + +```typescript +const documents = [ + "Document 1 text...", + "Document 2 text...", + "Document 3 text...", +]; + +const embeddingsArray = await embeddings.embedDocuments(documents); +console.log(embeddingsArray); +``` + +## Customizing Requests + +You can customize the base URL the SDK sends requests to by passing a `configuration` parameter: + +```typescript +const customEmbeddings = new DeepInfraEmbeddings({ + apiToken: "YOUR_API_TOKEN", + configuration: { + baseURL: "https://your_custom_url.com", + }, +}); +``` + +This allows you to route requests through a custom endpoint if needed. + +## Error Handling + +If the API token is not provided and cannot be found in the environment variables, an error will be thrown: + +```typescript +try { + const embeddings = new DeepInfraEmbeddings(); +} catch (error) { + console.error("DeepInfra API token not found"); +} +``` + +## Example + +Here’s a complete example of how to set up and use the `DeepInfraEmbeddings` class: + +```typescript +import { DeepInfraEmbeddings } from "@langchain/community/embeddings/deepinfra"; + +const embeddings = new DeepInfraEmbeddings({ + apiToken: "YOUR_API_TOKEN", + modelName: "sentence-transformers/clip-ViT-B-32", + batchSize: 512, +}); + +async function runExample() { + const queryEmbedding = await embeddings.embedQuery("Example query text."); + console.log("Query Embedding:", queryEmbedding); + + const documents = ["Text 1", "Text 2", "Text 3"]; + const documentEmbeddings = await embeddings.embedDocuments(documents); + console.log("Document Embeddings:", documentEmbeddings); +} + +runExample(); +``` + +## Feedback and Support + +For feedback or questions, please contact [feedback@deepinfra.com](mailto:feedback@deepinfra.com). diff --git a/examples/src/embeddings/deepinfra.ts b/examples/src/embeddings/deepinfra.ts new file mode 100644 index 000000000000..ed48eaf6a3b8 --- /dev/null +++ b/examples/src/embeddings/deepinfra.ts @@ -0,0 +1,12 @@ +import { DeepInfraEmbeddings } from "@langchain/community/embeddings/deepinfra"; + +const model = new DeepInfraEmbeddings({ + apiToken: process.env.DEEPINFRA_API_TOKEN, + batchSize: 1024, // Default value + modelName: "sentence-transformers/clip-ViT-B-32", // Default value +}); + +const embeddings = await model.embedQuery( + "Tell me a story about a dragon and a princess." +); +console.log(embeddings); diff --git a/examples/src/models/embeddings/deepinfra.ts b/examples/src/models/embeddings/deepinfra.ts new file mode 100644 index 000000000000..ed48eaf6a3b8 --- /dev/null +++ b/examples/src/models/embeddings/deepinfra.ts @@ -0,0 +1,12 @@ +import { DeepInfraEmbeddings } from "@langchain/community/embeddings/deepinfra"; + +const model = new DeepInfraEmbeddings({ + apiToken: process.env.DEEPINFRA_API_TOKEN, + batchSize: 1024, // Default value + modelName: "sentence-transformers/clip-ViT-B-32", // Default value +}); + +const embeddings = await model.embedQuery( + "Tell me a story about a dragon and a princess." +); +console.log(embeddings); diff --git a/libs/langchain-community/.gitignore b/libs/langchain-community/.gitignore index f88c0d76086e..3932edf88b61 100644 --- a/libs/langchain-community/.gitignore +++ b/libs/langchain-community/.gitignore @@ -142,6 +142,10 @@ embeddings/cohere.cjs embeddings/cohere.js embeddings/cohere.d.ts embeddings/cohere.d.cts +embeddings/deepinfra.cjs +embeddings/deepinfra.js +embeddings/deepinfra.d.ts +embeddings/deepinfra.d.cts embeddings/fireworks.cjs embeddings/fireworks.js embeddings/fireworks.d.ts diff --git a/libs/langchain-community/langchain.config.js b/libs/langchain-community/langchain.config.js index 3809415a8897..12b94735c697 100644 --- a/libs/langchain-community/langchain.config.js +++ b/libs/langchain-community/langchain.config.js @@ -69,6 +69,7 @@ export const config = { "embeddings/bedrock": "embeddings/bedrock", "embeddings/cloudflare_workersai": "embeddings/cloudflare_workersai", "embeddings/cohere": "embeddings/cohere", + "embeddings/deepinfra": "embeddings/deepinfra", "embeddings/fireworks": "embeddings/fireworks", "embeddings/googlepalm": "embeddings/googlepalm", "embeddings/googlevertexai": "embeddings/googlevertexai", diff --git a/libs/langchain-community/package.json b/libs/langchain-community/package.json index bbfaab6278b3..bb7a4876f2d2 100644 --- a/libs/langchain-community/package.json +++ b/libs/langchain-community/package.json @@ -1022,6 +1022,15 @@ "import": "./embeddings/cohere.js", "require": "./embeddings/cohere.cjs" }, + "./embeddings/deepinfra": { + "types": { + "import": "./embeddings/deepinfra.d.ts", + "require": "./embeddings/deepinfra.d.cts", + "default": "./embeddings/deepinfra.d.ts" + }, + "import": "./embeddings/deepinfra.js", + "require": "./embeddings/deepinfra.cjs" + }, "./embeddings/fireworks": { "types": { "import": "./embeddings/fireworks.d.ts", @@ -3096,6 +3105,10 @@ "embeddings/cohere.js", "embeddings/cohere.d.ts", "embeddings/cohere.d.cts", + "embeddings/deepinfra.cjs", + "embeddings/deepinfra.js", + "embeddings/deepinfra.d.ts", + "embeddings/deepinfra.d.cts", "embeddings/fireworks.cjs", "embeddings/fireworks.js", "embeddings/fireworks.d.ts", diff --git a/libs/langchain-community/src/embeddings/deepinfra.ts b/libs/langchain-community/src/embeddings/deepinfra.ts new file mode 100644 index 000000000000..6963902351a9 --- /dev/null +++ b/libs/langchain-community/src/embeddings/deepinfra.ts @@ -0,0 +1,181 @@ +import { getEnvironmentVariable } from "@langchain/core/utils/env"; +import { Embeddings, EmbeddingsParams } from "@langchain/core/embeddings"; +import { chunkArray } from "@langchain/core/utils/chunk_array"; + +/** + * The default model name to use for generating embeddings. + */ +const DEFAULT_MODEL_NAME = "sentence-transformers/clip-ViT-B-32"; + +/** + * The default batch size to use for generating embeddings. + * This is limited by the DeepInfra API to a maximum of 1024. + */ +const DEFAULT_BATCH_SIZE = 1024; + +/** + * Environment variable name for the DeepInfra API token. + */ +const API_TOKEN_ENV_VAR = "DEEPINFRA_API_TOKEN"; + +export interface DeepInfraEmbeddingsRequest { + inputs: string[]; + normalize?: boolean; + image?: string; + webhook?: string; +} + +/** + * Input parameters for the DeepInfra embeddings + */ +export interface DeepInfraEmbeddingsParams extends EmbeddingsParams { + /** + * The API token to use for authentication. + * If not provided, it will be read from the `DEEPINFRA_API_TOKEN` environment variable. + */ + apiToken?: string; + + /** + * The model ID to use for generating completions. + * Default: `sentence-transformers/clip-ViT-B-32` + */ + modelName?: string; + + /** + * The maximum number of texts to embed in a single request. This is + * limited by the DeepInfra API to a maximum of 1024. + */ + batchSize?: number; +} + +/** + * Response from the DeepInfra embeddings API. + */ +export interface DeepInfraEmbeddingsResponse { + /** + * The embeddings generated for the input texts. + */ + embeddings: number[][]; + /** + * The number of tokens in the input texts. + */ + input_tokens: number; + /** + * The status of the inference. + */ + request_id?: string; +} + +/** + * A class for generating embeddings using the DeepInfra API. + * @example + * ```typescript + * // Embed a query using the DeepInfraEmbeddings class + * const model = new DeepInfraEmbeddings(); + * const res = await model.embedQuery( + * "What would be a good company name for a company that makes colorful socks?", + * ); + * console.log({ res }); + * ``` + */ +export class DeepInfraEmbeddings + extends Embeddings + implements DeepInfraEmbeddingsParams +{ + apiToken: string; + + batchSize: number; + + modelName: string; + + /** + * Constructor for the DeepInfraEmbeddings class. + * @param fields - An optional object with properties to configure the instance. + */ + constructor( + fields?: Partial & { + verbose?: boolean; + } + ) { + const fieldsWithDefaults = { + modelName: DEFAULT_MODEL_NAME, + batchSize: DEFAULT_BATCH_SIZE, + ...fields, + }; + + super(fieldsWithDefaults); + + const apiKey = + fieldsWithDefaults?.apiToken || getEnvironmentVariable(API_TOKEN_ENV_VAR); + + if (!apiKey) { + throw new Error("DeepInfra API token not found"); + } + + this.modelName = fieldsWithDefaults?.modelName ?? this.modelName; + this.batchSize = fieldsWithDefaults?.batchSize ?? this.batchSize; + this.apiToken = apiKey; + } + + /** + * Generates embeddings for an array of texts. + * @param inputs - An array of strings to generate embeddings for. + * @returns A Promise that resolves to an array of embeddings. + */ + async embedDocuments(inputs: string[]): Promise { + const batches = chunkArray(inputs, this.batchSize); + + const batchRequests = batches.map((batch: string[]) => + this.embeddingWithRetry({ + inputs: batch, + }) + ); + + const batchResponses = await Promise.all(batchRequests); + + const out: number[][] = []; + + for (let i = 0; i < batchResponses.length; i += 1) { + const batch = batches[i]; + const { embeddings } = batchResponses[i]; + for (let j = 0; j < batch.length; j += 1) { + out.push(embeddings[j]); + } + } + + return out; + } + + /** + * Generates an embedding for a single text. + * @param text - A string to generate an embedding for. + * @returns A Promise that resolves to an array of numbers representing the embedding. + */ + async embedQuery(text: string): Promise { + const { embeddings } = await this.embeddingWithRetry({ + inputs: [text], + }); + return embeddings[0]; + } + + /** + * Generates embeddings with retry capabilities. + * @param request - An object containing the request parameters for generating embeddings. + * @returns A Promise that resolves to the API response. + */ + private async embeddingWithRetry( + request: DeepInfraEmbeddingsRequest + ): Promise { + const response = await this.caller.call(() => + fetch(`https://api.deepinfra.com/v1/inference/${this.modelName}`, { + method: "POST", + headers: { + Authorization: `Bearer ${this.apiToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(request), + }).then((res) => res.json()) + ); + return response as DeepInfraEmbeddingsResponse; + } +} diff --git a/libs/langchain-community/src/embeddings/tests/deepinfra.int.test.ts b/libs/langchain-community/src/embeddings/tests/deepinfra.int.test.ts new file mode 100644 index 000000000000..748bc1dc6dc9 --- /dev/null +++ b/libs/langchain-community/src/embeddings/tests/deepinfra.int.test.ts @@ -0,0 +1,34 @@ +import { test, expect } from "@jest/globals"; +import { DeepInfraEmbeddings } from "../deepinfra.js"; + +test("Test DeepInfraEmbeddings.embedQuery", async () => { + const embeddings = new DeepInfraEmbeddings(); + const res = await embeddings.embedQuery("Hello world"); + expect(typeof res[0]).toBe("number"); +}); + +test("Test DeepInfraEmbeddings.embedDocuments", async () => { + const embeddings = new DeepInfraEmbeddings(); + const res = await embeddings.embedDocuments(["Hello world", "Bye bye"]); + expect(res).toHaveLength(2); + expect(typeof res[0][0]).toBe("number"); + expect(typeof res[1][0]).toBe("number"); +}); + +test("Test DeepInfraEmbeddings concurrency", async () => { + const embeddings = new DeepInfraEmbeddings({ + batchSize: 1, + }); + const res = await embeddings.embedDocuments([ + "Hello world", + "Bye bye", + "we need", + "at least", + "six documents", + "to test concurrency", + ]); + expect(res).toHaveLength(6); + expect(res.find((embedding) => typeof embedding[0] !== "number")).toBe( + undefined + ); +});