IndexCollectionVectors.java

package io.anserini.ann;

import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Paths;

import com.robrua.nlp.bert.Bert;
import io.anserini.index.IndexArgs;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
/**
* create a vectors file on index contents using easyBERT
*/
public class IndexCollectionVectors {

  public static void main(String[] args) throws Exception {
    Bert bert = Bert.load(Paths.get("/path/to/bert"));
    Directory dir = FSDirectory.open(Paths.get("/path/to/lucene-index"));
    DirectoryReader reader = DirectoryReader.open(dir);
    FileOutputStream fos = new FileOutputStream("vectors.txt");
    for (int i = 0; i < reader.maxDoc(); i++) {{
      Document d = reader.document(i);
      float[] vector = bert.embedSequence(d.get(IndexArgs.CONTENTS));
      String id = d.get("id");
      fos.write(id.getBytes());
      for (float f : vector) {
        String s = " " + f;
        fos.write(s.getBytes());
      }
      fos.write("\n".getBytes());
    }
    fos.flush();
    fos.close();
  }
}