-
Notifications
You must be signed in to change notification settings - Fork 801
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Embeddings][OpenAI] Support embeddings via engine.embeddings.create()
- Loading branch information
1 parent
66dd646
commit ef27ebc
Showing
17 changed files
with
1,112 additions
and
31 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,14 @@ | ||
| # WebLLM Get Started App | ||
|
|
||
| This folder provides a minimum demo to show WebLLM API in a webapp setting. | ||
| To try it out, you can do the following steps under this folder | ||
|
|
||
| ```bash | ||
| npm install | ||
| npm start | ||
| ``` | ||
|
|
||
| Note if you would like to hack WebLLM core package. | ||
| You can change web-llm dependencies as `"file:../.."`, and follow the build from source | ||
| instruction in the project to build webllm locally. This option is only recommended | ||
| if you would like to hack WebLLM core package. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| { | ||
| "name": "embeddings-example", | ||
| "version": "0.1.0", | ||
| "private": true, | ||
| "scripts": { | ||
| "start": "parcel src/embeddings.html --port 8885", | ||
| "build": "parcel build src/embeddings.html --dist-dir lib" | ||
| }, | ||
| "devDependencies": { | ||
| "buffer": "^5.7.1", | ||
| "parcel": "^2.8.3", | ||
| "process": "^0.11.10", | ||
| "tslib": "^2.3.1", | ||
| "typescript": "^4.9.5", | ||
| "url": "^0.11.3" | ||
| }, | ||
| "dependencies": { | ||
| "@mlc-ai/web-llm": "file:../..", | ||
| "langchain": "0.2.15" | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,23 @@ | ||
| <!doctype html> | ||
| <html> | ||
| <script> | ||
| webLLMGlobal = {}; | ||
| </script> | ||
| <body> | ||
| <h2>WebLLM Test Page</h2> | ||
| Open console to see output | ||
| <br /> | ||
| <br /> | ||
| <label id="init-label"> </label> | ||
|
|
||
| <h3>Prompt</h3> | ||
| <label id="prompt-label"> </label> | ||
|
|
||
| <h3>Response</h3> | ||
| <label id="generate-label"> </label> | ||
| <br /> | ||
| <label id="stats-label"> </label> | ||
|
|
||
| <script type="module" src="./embeddings.ts"></script> | ||
| </body> | ||
| </html> |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,147 @@ | ||
| import * as webllm from "@mlc-ai/web-llm"; | ||
| import { MemoryVectorStore } from "langchain/vectorstores/memory"; | ||
| import type { EmbeddingsInterface } from "@langchain/core/embeddings"; | ||
| import type { Document } from "@langchain/core/documents"; | ||
|
|
||
| function setLabel(id: string, text: string) { | ||
| const label = document.getElementById(id); | ||
| if (label == null) { | ||
| throw Error("Cannot find label " + id); | ||
| } | ||
| label.innerText = text; | ||
| } | ||
|
|
||
| const initProgressCallback = (report: webllm.InitProgressReport) => { | ||
| setLabel("init-label", report.text); | ||
| }; | ||
|
|
||
| // For integration with Langchain | ||
| class WebLLMEmbeddings implements EmbeddingsInterface { | ||
| engine: webllm.MLCEngineInterface; | ||
| constructor(engine: webllm.MLCEngineInterface) { | ||
| this.engine = engine; | ||
| } | ||
|
|
||
| async _embed(texts: string[]): Promise<number[][]> { | ||
| const reply = await this.engine.embeddings.create({ input: texts }); | ||
| const result: number[][] = []; | ||
| for (let i = 0; i < texts.length; i++) { | ||
| result.push(reply.data[i].embedding); | ||
| } | ||
| return result; | ||
| } | ||
|
|
||
| async embedQuery(document: string): Promise<number[]> { | ||
| return this._embed([document]).then((embeddings) => embeddings[0]); | ||
| } | ||
|
|
||
| async embedDocuments(documents: string[]): Promise<number[][]> { | ||
| return this._embed(documents); | ||
| } | ||
| } | ||
|
|
||
| // Prepare inputs | ||
| const documents_og = ["The Data Cloud!", "Mexico City of Course!"]; | ||
| const queries_og = ["what is snowflake?", "Where can I get the best tacos?"]; | ||
| const documents: string[] = []; | ||
| const queries: string[] = []; | ||
| const query_prefix = | ||
| "Represent this sentence for searching relevant passages: "; | ||
| // Process according to Snowflake model | ||
| documents_og.forEach(function (item, index) { | ||
| documents[index] = `[CLS] ${item} [SEP]`; | ||
| }); | ||
| queries_og.forEach(function (item, index) { | ||
| queries[index] = `[CLS] ${query_prefix}${item} [SEP]`; | ||
| }); | ||
| console.log("Formatted documents: ", documents); | ||
| console.log("Formatted queries: ", queries); | ||
|
|
||
| // Using webllm's API | ||
| async function webllmAPI() { | ||
| // b4 means the max batch size is compiled as 4. That is, the model can process 4 inputs in a | ||
| // batch. If given more than 4, the model will forward multiple times. The larger the max batch | ||
| // size, the more memory it consumes. | ||
| // const selectedModel = "snowflake-arctic-embed-m-q0f32-MLC-b32"; | ||
| const selectedModel = "snowflake-arctic-embed-m-q0f32-MLC-b4"; | ||
| const engine: webllm.MLCEngineInterface = await webllm.CreateMLCEngine( | ||
| selectedModel, | ||
| { | ||
| initProgressCallback: initProgressCallback, | ||
| logLevel: "INFO", // specify the log level | ||
| }, | ||
| ); | ||
|
|
||
| const docReply = await engine.embeddings.create({ input: documents }); | ||
| console.log(docReply); | ||
| console.log(docReply.usage); | ||
|
|
||
| const queryReply = await engine.embeddings.create({ input: queries }); | ||
| console.log(queryReply); | ||
| console.log(queryReply.usage); | ||
|
|
||
| // Calculate similarity (we use langchain here, but any method works) | ||
| const vectorStore = await MemoryVectorStore.fromExistingIndex( | ||
| new WebLLMEmbeddings(engine), | ||
| ); | ||
| // See score | ||
| for (let i = 0; i < queries_og.length; i++) { | ||
| console.log(`Similarity with: ${queries_og[i]}`); | ||
| for (let j = 0; j < documents_og.length; j++) { | ||
| const similarity = vectorStore.similarity( | ||
| queryReply.data[i].embedding, | ||
| docReply.data[j].embedding, | ||
| ); | ||
| console.log(`${documents_og[j]}: ${similarity}`); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Alternatively, integrating with Langchain's API | ||
| async function langchainAPI() { | ||
| // b4 means the max batch size is compiled as 4. That is, the model can process 4 inputs in a | ||
| // batch. If given more than 4, the model will forward multiple times. The larger the max batch | ||
| // size, the more memory it consumes. | ||
| // const selectedModel = "snowflake-arctic-embed-m-q0f32-MLC-b32"; | ||
| const selectedModel = "snowflake-arctic-embed-m-q0f32-MLC-b4"; | ||
| const engine: webllm.MLCEngineInterface = await webllm.CreateMLCEngine( | ||
| selectedModel, | ||
| { | ||
| initProgressCallback: initProgressCallback, | ||
| logLevel: "INFO", // specify the log level | ||
| }, | ||
| ); | ||
|
|
||
| const vectorStore = await MemoryVectorStore.fromExistingIndex( | ||
| new WebLLMEmbeddings(engine), | ||
| ); | ||
| const document0: Document = { | ||
| pageContent: documents[0], | ||
| metadata: {}, | ||
| }; | ||
| const document1: Document = { | ||
| pageContent: documents[1], | ||
| metadata: {}, | ||
| }; | ||
| await vectorStore.addDocuments([document0, document1]); | ||
|
|
||
| const similaritySearchResults0 = await vectorStore.similaritySearch( | ||
| queries[0], | ||
| 1, | ||
| ); | ||
| for (const doc of similaritySearchResults0) { | ||
| console.log(`* ${doc.pageContent}`); | ||
| } | ||
|
|
||
| const similaritySearchResults1 = await vectorStore.similaritySearch( | ||
| queries[1], | ||
| 1, | ||
| ); | ||
| for (const doc of similaritySearchResults1) { | ||
| console.log(`* ${doc.pageContent}`); | ||
| } | ||
| } | ||
|
|
||
| // Select one to run | ||
| webllmAPI(); | ||
| // langchainAPI(); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.