/
math.go
65 lines (55 loc) · 1.09 KB
/
math.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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
package localdb
import (
"math"
"github.com/byebyebruce/chat2code/vectordb"
)
type Float interface {
~float32 | ~float64
}
// Cosine 计算cosine [-1,1],越大越相似
func Cosine[T Float](a []T, b []T) T {
count := 0
length_a := len(a)
length_b := len(b)
if length_a > length_b {
count = length_a
} else {
count = length_b
}
sumA := 0.0
s1 := 0.0
s2 := 0.0
for k := 0; k < count; k++ {
if k >= length_a {
s2 += math.Pow(float64(b[k]), 2)
continue
}
if k >= length_b {
s1 += math.Pow(float64(a[k]), 2)
continue
}
sumA += float64(a[k]) * float64(b[k])
s1 += math.Pow(float64(a[k]), 2)
s2 += math.Pow(float64(b[k]), 2)
}
return T(sumA / (math.Sqrt(s1) * math.Sqrt(s2)))
}
type bigHeap []*vectordb.Match
func (b bigHeap) Len() int {
return len(b)
}
func (b bigHeap) Less(i, j int) bool {
return b[i].Score > b[j].Score
}
func (b bigHeap) Swap(i, j int) {
b[i], b[j] = b[j], b[i]
}
func (b *bigHeap) Push(x interface{}) {
*b = append(*b, x.(*vectordb.Match))
}
func (b *bigHeap) Pop() interface{} {
l := len(*b) - 1
a := (*b)[l]
*b = (*b)[:l]
return a
}