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!