bench_hnsw.go

//go:build ignore
// +build ignore

package main

import (
	"log"
	"net/http"
	_ "net/http/pprof"

	"github.com/habedi/hann/core"
	"github.com/habedi/hann/example"
	"github.com/habedi/hann/hnsw"
)

func main() {
	// Start the pprof HTTP server on port 6060. It exposes the profiling
	// endpoints at /debug/pprof/.
	go func() {
		log.Println("Starting pprof server on :6060")
		if err := http.ListenAndServe("localhost:6060", nil); err != nil {
			log.Printf("pprof server failed: %v", err)
		}
	}()

	// Benchmark the HNSW index on the FashionMNIST and Glove datasets.
	BenchHNSWIndexFashionMNIST()
	BenchHNSWIndexGlove25()
	BenchHNSWIndexGlove200()
}

func BenchHNSWIndexFashionMNIST() {
	factory := func() core.Index {
		dimension := 784
		M := 32
		ef := 300
		index, err := hnsw.New(dimension, hnsw.WithM(M), hnsw.WithEf(ef))
		if err != nil {
			log.Fatalf("Failed to create HNSW index: %v", err)
		}
		return index
	}

	example.RunDataset(factory, "fashion-mnist-784-euclidean",
		"example/data/nearest-neighbors-datasets", 100, -1, 5)
}

func BenchHNSWIndexGlove25() {
	factory := func() core.Index {
		dimension := 25
		index, err := hnsw.New(dimension, hnsw.WithMetric(core.Cosine))
		if err != nil {
			log.Fatalf("Failed to create HNSW index: %v", err)
		}
		return index
	}

	example.RunDataset(factory, "glove-25-angular",
		"example/data/nearest-neighbors-datasets", 100, -1, 5)
}

func BenchHNSWIndexGlove200() {
	factory := func() core.Index {
		dimension := 200
		index, err := hnsw.New(dimension, hnsw.WithMetric(core.Cosine))
		if err != nil {
			log.Fatalf("Failed to create HNSW index: %v", err)
		}
		return index
	}

	example.RunDataset(factory, "glove-200-angular",
		"example/data/nearest-neighbors-datasets", 100, -1, 5)
}