Pinecone + OpenAI Embedding Vector Search

This is a simple example app that shows how to use the Pinecone vector database
to easily search for words and phrases matching an input query. We use
the OpenAI text-embedding-ada-002
model to generate embedding vectors for us.
Note that for this demo to work, your Pinecone database must have the Metric
set to cosine
(cosine similarity) and the Dimensions
set to 1,536 to work
properly with OpenAI's text-embedding-ada-002
model.
Populating your vector database
First use the OpenAI API to generate embedding
vectors for a list of animal names in animals.json
. This is in
functions/vectors/populate.js
. Note that we batch requests to save on
processing time. To insert over 300 animals it takes about 10s. You could
parallelize both the embedding creations and upserts to speed this up using
e.g. Promise.all([...])
.
// This file will populate your Pinecone vector databse with
// a list of animal names
const lib = require('lib')({token: process.env.STDLIB_SECRET_TOKEN});
const animals = require('../../reference/animals.json');
let startTime = new Date().valueOf();
// First, create embeddings for animals
// We'll batch create embedding vectors, reference:
// https://autocode.com/openai/snippets/cachsnpt_8X8A56WqhFZ1dZe67fE8wDG7JDv7qgKc72Sf/
let inputs = animals.slice();
let embeddings = [];
// batch inputs so we don’t exceed token limits
// tokens aren’t exactly words, so we’ll limit tokenCount to 4096
// in case of weird characters, this should handle most input variatons
while (inputs.length) {
let tokenCount = 0;
let batchedInputs = [];
while (inputs.length && tokenCount < 4096) {
let input = inputs.shift();
batchedInputs.push(input);
tokenCount += input.split(' ').length;
}
let embeddingResult = await lib.openai.playground['@0.0.4'].embeddings.create({
model: `text-embedding-ada-002`,
input: batchedInputs
});
embeddings = embeddings.concat(
embeddingResult.data.map(entry => entry.embedding)
);
}
let creationTime = new Date().valueOf() - startTime;
// We'll prepare the data for insertion into Pinecone
let vectors = animals.map((animal, i) => {
return {
id: animal,
metadata: {
name: animal
},
values: embeddings[i]
}
});
// Now we'll add our vectors into Pinecone
// We can only do about 250 at a time due to size limitations
let insertBatches = [];
while (vectors.length) {
let batchedVectors = vectors.splice(0, 250);
let pineconeResult = await lib.pinecone.index['@0.0.1'].vectors.upsert({
namespace: `animals`,
vectors: batchedVectors
});
insertBatches.push(pineconeResult);
}
let totalTime = new Date().valueOf() - startTime;
// Return time and batch stats
return {
creationTime,
upsertTime: totalTime - creationTime,
totalTime,
insertBatches
};
Querying your vector Database
Querying is easy, just generate an embedding vector based on user input
and make a query request to Pinecone. Not that plain english tends to work
better to get the results you want, e.g. to query for mammals animals that are mammals
works a little better than mammals
. This is an artifact of the embedding vectors
created by text-embedding-ada-002
, other models may perform differently.
In our experience, the more context you give the model the more accurate the
comparisons.
// Query animal names based on input
const lib = require('lib')({token: process.env.STDLIB_SECRET_TOKEN});
let startTime = new Date().valueOf();
// Input could be, e.g. "animals that are mammals"
// Query works better if plain english
let query = context.params.query;
// Convert to an embedding vector
let embeddingResult = await lib.openai.playground['@0.0.4'].embeddings.create({
model: `text-embedding-ada-002`,
input: [query]
});
let vector = embeddingResult.data[0].embedding;
let creationTime = new Date().valueOf() - startTime;
// Query pinecone for result
let pineconeResult = await lib.pinecone.index['@0.0.1'].query({
namespace: `animals`,
vector: vector,
topK: 5,
includeMetadata: true,
includeValues: false
});
let matches = pineconeResult.matches;
let totalTime = new Date().valueOf() - startTime;
// Return query matches
return {
creationTime,
queryTime: totalTime - creationTime,
totalTime,
query,
matches
};
That's all!
This is a very simple example of how you can use embedding vectors
and the Pinecone vector database to search for related words and phrases.
It makes querying against large datasets and building custom search with
a high degree of accuracy easy to implement on your own instead of relying on
tools like Algolia.
For help with Autocode, join the Autocode Discord server at
https://discord.gg/autocode - enjoy!