-
Notifications
You must be signed in to change notification settings - Fork 4
/
quantization_f16_test.go
51 lines (47 loc) · 1.22 KB
/
quantization_f16_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
package bbq
import "testing"
func TestFloat16Quantization(t *testing.T) {
vecs := NewRandVectorSet(1000, *dim, nil)
mem := NewMemoryBackend(*dim)
quant := NewQuantizedMemoryBackend(*dim, Float16Quantization{})
for i, v := range vecs {
mem.PutVector(ID(i), v)
quant.PutVector(ID(i), v)
}
target := NewRandVector(*dim, nil)
memrs, err := FullTableScanSearch(mem, target, 20)
if err != nil {
t.Fatal(err)
}
qrs, err := FullTableScanSearch(quant, target, 20)
if err != nil {
t.Fatal(err)
}
recall := memrs.ComputeRecall(qrs, 10)
t.Logf("Recall %0.4f\n", recall)
t.Logf("\n%s\n%s", memrs, qrs)
}
func TestFloat16Backend(t *testing.T) {
vecs := NewRandVectorSet(1000, *dim, nil)
quant := NewQuantizedMemoryBackend(*dim, Float16Quantization{})
store, err := NewVectorStore(quant, *nBasis, WithPrespill(2))
if err != nil {
t.Fatal(err)
}
err = store.AddVectorsWithOffset(0, vecs)
if err != nil {
t.Fatal(err)
}
target := NewRandVector(*dim, nil)
qrs, err := FullTableScanSearch(quant, target, 20)
if err != nil {
t.Fatal(err)
}
rs, err := store.FindNearest(target, 20, 20000, 2)
if err != nil {
t.Fatal(err)
}
recall := rs.ComputeRecall(qrs, 10)
t.Logf("Recall %0.4f\n", recall)
t.Logf("\n%s\n%s", rs, qrs)
}