Improve document embedding speed 10x

This commit is contained in:
timothycarambat
2023-07-26 18:22:03 -07:00
parent 904ba56ad3
commit 79ac2b7385
3 changed files with 33 additions and 24 deletions
+13
View File
@@ -17,6 +17,19 @@ class OpenAi {
? data[0].embedding
: null;
}
async embedTextChunks(chunks = []) {
const {
data: { data },
} = await this.openai.createEmbedding({
model: "text-embedding-ada-002",
input: chunks,
});
return data.length > 0 &&
data.every((embd) => embd.hasOwnProperty("embedding"))
? data.map((embd) => embd.embedding)
: null;
}
}
module.exports = {
@@ -135,7 +135,7 @@ class Chroma {
const documentVectors = [];
const cacheInfo = [];
const vectors = [];
const vectorValues = await openai.embedTextChunks(textChunks);
const submission = {
ids: [],
embeddings: [],
@@ -143,23 +143,21 @@ class Chroma {
documents: [],
};
for (const textChunk of textChunks) {
const vectorValues = await openai.embedTextChunk(textChunk);
if (!!vectorValues) {
if (!!vectorValues && vectorValues.length > 0) {
for (const [i, vector] of vectorValues.entries()) {
const vectorRecord = {
id: v4(),
values: vectorValues,
values: vector,
// [DO NOT REMOVE]
// LangChain will be unable to find your text if you embed manually and dont include the `text` key.
// https://github.com/hwchase17/langchainjs/blob/2def486af734c0ca87285a48f1a04c057ab74bdf/langchain/src/vectorstores/pinecone.ts#L64
metadata: { ...metadata, text: textChunk },
metadata: { ...metadata, text: textChunks[i] },
};
submission.ids.push(vectorRecord.id);
submission.embeddings.push(vectorRecord.values);
submission.metadatas.push(metadata);
submission.documents.push(textChunk);
submission.documents.push(textChunks[i]);
vectors.push(vectorRecord);
documentVectors.push({
@@ -172,11 +170,11 @@ class Chroma {
values: vectorValues,
metadata: vectorRecord.metadata,
});
} else {
console.error(
"Could not use OpenAI to embed document chunk! This document will not be recorded."
);
}
} else {
console.error(
"Could not use OpenAI to embed document chunk! This document will not be recorded."
);
}
const { client } = await this.connect();
@@ -162,7 +162,7 @@ class Pinecone {
const documentVectors = [];
const cacheInfo = [];
const vectors = [];
const vectorValues = await openai.embedTextChunks(textChunks);
const submission = {
ids: [],
embeddings: [],
@@ -170,23 +170,21 @@ class Pinecone {
documents: [],
};
for (const textChunk of textChunks) {
const vectorValues = await openai.embedTextChunk(textChunk);
if (!!vectorValues) {
if (!!vectorValues && vectorValues.length > 0) {
for (const [i, vector] of vectorValues.entries()) {
const vectorRecord = {
id: v4(),
values: vectorValues,
values: vector,
// [DO NOT REMOVE]
// LangChain will be unable to find your text if you embed manually and dont include the `text` key.
// https://github.com/hwchase17/langchainjs/blob/2def486af734c0ca87285a48f1a04c057ab74bdf/langchain/src/vectorstores/pinecone.ts#L64
metadata: { ...metadata, text: textChunk },
metadata: { ...metadata, text: textChunks[i] },
};
submission.ids.push(vectorRecord.id);
submission.embeddings.push(vectorRecord.values);
submission.metadatas.push(metadata);
submission.documents.push(textChunk);
submission.documents.push(textChunks[i]);
vectors.push(vectorRecord);
documentVectors.push({
@@ -199,11 +197,11 @@ class Pinecone {
values: vectorValues,
metadata: vectorRecord.metadata,
});
} else {
console.error(
"Could not use OpenAI to embed document chunk! This document will not be recorded."
);
}
} else {
console.error(
"Could not use OpenAI to embed document chunk! This document will not be recorded."
);
}
if (vectors.length > 0) {