Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

community[minor]: DeepInfra embeddings integration #1 #5382

Merged
merged 14 commits into from
May 31, 2024
130 changes: 130 additions & 0 deletions docs/core_docs/docs/integrations/text_embedding/deepinfra.mdx
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
---
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";

<IntegrationInstallTooltip></IntegrationInstallTooltip>

```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).


12 changes: 12 additions & 0 deletions examples/src/embeddings/deepinfra.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
import { DeepInfraEmbeddings } from "@langchain/community/embeddings/deepinfra";
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work on the PR! I've flagged the addition of environment variable access using process.env for the maintainers to review. Keep up the good work!


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);
12 changes: 12 additions & 0 deletions examples/src/models/embeddings/deepinfra.ts
Original file line number Diff line number Diff line change
@@ -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);
4 changes: 4 additions & 0 deletions libs/langchain-community/.gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions libs/langchain-community/langchain.config.js
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
13 changes: 13 additions & 0 deletions libs/langchain-community/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -1022,6 +1022,15 @@
"import": "./embeddings/cohere.js",
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey there! I noticed that this PR introduces a new dependency for "./embeddings/deepinfra" in the package.json file. This change may impact the project's dependencies, so I'm flagging it for the maintainers to review. Keep up the great work!

"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",
Expand Down Expand Up @@ -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",
Expand Down
181 changes: 181 additions & 0 deletions libs/langchain-community/src/embeddings/deepinfra.ts
Original file line number Diff line number Diff line change
@@ -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<DeepInfraEmbeddingsParams> & {
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<number[][]> {
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<number[]> {
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<DeepInfraEmbeddingsResponse> {
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;
}
}
Loading
Loading