Balltree data structure implemented in ocaml-torch. Currently only supports Euclidean distance, and takes more memory than the scikit-learn balltree.
Also has an example webserver application, built on top ("balltree as a service").
Planned:
[] more/ arbitrary distances
[] serialization to disk
To build and run the example container:
sudo docker build -t balltree .
sudo docker run -p 3000:3000 balltree
(Takes 4+GB of RAM).
$ curl http://localhost:3000/getnn/Right%20noble%20Burgundy,/10
{"query":"Right noble Burgundy,","n_neighbours":10,"distances":[0.0,118.38802337646484,122.64796447753906,124.86466217041016,125.55013275146484,131.27127075195312,134.02183532714844,135.88157653808594,136.78036499023438,136.79348754882812],"indices":[106,267,66,701,428,20,270,599,779,728],"matched_lines":["Right noble Burgundy,\r","By Providence divine.\r","another.\r","CLOWN, Servant to Othello\r","IAGO, his Ancient\r"," his BRETHREN\r","Tis true.\r"," Who, Hero?\r","Alas, Iago,\r","By whom, Aeneas?\r"]}You could alternatively use data similar to input_sentences.txt by:
wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt works_of_shakespeare.txt | sort | uniq | shuf | head -n 1000 > input_sentences.txt
We only keep 1000 sentences, since the (distil)bert embeddings are still high dimensional.
$ du -sh sentence_embeddings.txt
574M sentence_embeddings.txt