ベクトルデータベース Pinecornにベクトルを登録した #Node.js

はじめに

「Pinecone」は、フルマネージドなベクトルデータベースで、HTTPプロトコルによる
ベクトルの登録、検索を無料で試すごとができる。ベクトル検索用のサーバ本体、サーバ側のプログラムを自前で用意することなく、高速にベクトル検索できるのが魅力的だと思った。

環境

windows10
nodejs v16.14.0
VSCode

やったこと

類似画像検索てきなことを行ってみた。

PinecornにIndexを作成。また、APIキーも作成。
バッチプログラムで任意のフォルダ配下の複数のjpeg画像（約6万ファイル）を読み込みTensorflow.jsでベクトルを生成した。
生成されたベクトルを100個づつPinecornに登録した。

Pinecornに対する操作は、Indexの作成から行えるPinecornのライブラリがあり、使ってみたがベクトル登録時のエラーが解消されず使用をあきらめ、単純にfetch APIでPOSTリクエストを投げることにした。。。。

Pinecorn側に生成したIndexについて

画面から「Index名」、「登録するベクトルのディメンション」、「検索に使用する距離メトリクス」、「Pinecornが動作するPodType」を登録、選択するのみで、Indexがつくられる。また、簡単に削除することもできる。

知識不足でディメンションの値がよくわからなかったが、生成されたベクトルのサイズを実際に調べて1280だったので1280を設定した。また、検索はコサイン類似度を選択した。

Indexを生成すると、登録、検索など各操作のcurlコマンドのサンプルが画表示されるようになる。

作成したプログラムについて

####　登録
1回のPOSTリクエストで最大100個までベクトルを登録することが可能。下のvectorlistは、最大100個のjsonの配列。また、送信先のURLは、Pincecornの管理画面のcurlコマンドのサンプルからとってきた。

async function upsertVector(vectorlist){
    const res = await fetch(pineconeUrl,{
        method:'POST',
        headers : {
          'Content-Type': 'application/json',
          'Api-Key': apiKey
        },
        body : JSON.stringify({
          vectors: vectorlist,
          namespace: 'imageindex'
        })
      });
}

Jsonは、登録するベクトルそのものと、それに紐づく一意idの組み合わせ。idにはファイル名を使用した。また、やらなかったけどもメタデータも付加することできる。

function createVector(file,featureVector){
    const vector = {};
    
    vector.id = file;
    vector.values = featureVector;
    return vector;
}

【登録に使用したコード全文】

import * as fs from 'fs';
import * as tf from '@tensorflow/tfjs-node';
import fetch from "node-fetch";

const folderPath = "C:/temp/images";
const pineconeUrl = "https://imageindex-c5a7176.svc.us-east-1-aws.pinecone.io/vectors/upsert";
const apiKey = "<pinecornのAPIキー>";

async function getVector(model,file) {
    // Load the image as a tensor
    const imageBuffer = fs.readFileSync(folderPath + "/" + file);
    const imageTensor = tf.node.decodeImage(imageBuffer);
    const imageResized = tf.image.resizeBilinear(imageTensor, [96, 96]);
    const image = imageResized.expandDims(0).toFloat().div(tf.scalar(127)).sub(tf.scalar(1));

    const features =  model.predict(image);

    const float32array = features.dataSync();
    const featureVector = Array.from(float32array.slice(0));

    return featureVector;
}

function createVector(file,featureVector){
    const vector = {};
    
    vector.id = file;
    vector.values = featureVector;
    return vector;
}

async function upsertVector(vectorlist){
    const res = await fetch(pineconeUrl,{
        method:'POST',
        headers : {
          'Content-Type': 'application/json',
          'Api-Key': apiKey
        },
        body : JSON.stringify({
          vectors: vectorlist,
          namespace: 'imageindex'
        })
      });
}

async function main() {

    const model = await tf.loadGraphModel("https://tfhub.dev/google/tfjs-model/imagenet/mobilenet_v2_100_96/feature_vector/2/default/1", { fromTFHub: true });
    console.log('load model');

    let vectorlist = [];
    const filelist = fs.readdirSync(folderPath);
    for ( const file of filelist) {
        console.log(file);

        const featureVector = await getVector(model,file);
        vectorlist.push(createVector(file,featureVector));

        if (vectorlist.length == 100){
            try {
                await upsertVector(vectorlist);
            }catch (e){
                console.log(e);
            }finally {
                vectorlist = [];
            }
        }
    };
    
    // last batch
    if (vectorlist.length > 0){
        await upsertVector(vectorlist);
    }
}

main();

検索

実際に行いたいのはブラウザからの検索であるが、まずは動作確認のため、スクリプトから検索を試してみた。

登録にも使用した、画像ファイル(1011516521.jpg)から生成したベクトルを使用し、
それに近いベクトル5件を取得するリクエストを行ってみた。「includeMetadata」「includeValues」をそれぞれtrueに設定すると、メタデータ（登録していたら）とベクトルの値が戻り値に含まれる。
また、登録時にnamespaceを指定していたので、検索時も同じnamespaceを指定する。送信先のurlはQUERY用のURLを指定した。

    const res = await fetch(pineconeUrl,{
        method:'POST',
        headers : {
          'Content-Type': 'application/json',
          'Api-Key': apiKey
        },
        body: JSON.stringify({
    		vector: featureVector,
    		topK: 5,
    		includeMetadata: false,
    		includeValues: false,
    		namespace: 'imageindex'
  		})
      });

【検索プログラム全文】

import * as fs from 'fs';
import * as tf from '@tensorflow/tfjs-node';
import fetch from "node-fetch";

const folderPath = "C:/temp/images";
const pineconeUrl = "https://imageindex-c5a7176.svc.us-east-1-aws.pinecone.io/query";
const apiKey = "<pinecornのAPIキー>";

async function getVector(model,file) {
    // Load the image as a tensor
    const imageBuffer = fs.readFileSync(folderPath + "/" + file);
    const imageTensor = tf.node.decodeImage(imageBuffer);
    const imageResized = tf.image.resizeBilinear(imageTensor, [96, 96]);
    const image = imageResized.expandDims(0).toFloat().div(tf.scalar(127)).sub(tf.scalar(1));

    const features =  model.predict(image);

    // Get the: feature vector as a flat array
    const float32array = features.dataSync();
    const featureVector = Array.from(float32array.slice(0));

    return featureVector;
}

async function queryVector(featureVector){
    const res = await fetch(pineconeUrl,{
        method:'POST',
        headers : {
          'Content-Type': 'application/json',
          'Api-Key': apiKey
        },
        body: JSON.stringify({
    		vector: featureVector,
    		topK: 5,
    		includeMetadata: false,
    		includeValues: false,
    		namespace: 'imageindex'
  		})
      });

    const json = await res.json();
    return json;
}

async function main() {
    const model = await tf.loadGraphModel("https://tfhub.dev/google/tfjs-model/imagenet/mobilenet_v2_100_96/feature_vector/2/default/1", { fromTFHub: true });
    console.log('load model');

    const file = "1011516521.jpg";
    const featureVector = await getVector(model,file);

    const res = await queryVector(featureVector);
    console.log(res);
}

main();

検索結果

VSCodeから実行して、コンソールに出力するようにした。

おわりに

ブラウザからも検索を実行してみる。