156 lines
3.3 KiB
Go
156 lines
3.3 KiB
Go
package embedding
|
|
|
|
import (
|
|
"context"
|
|
"math"
|
|
"testing"
|
|
)
|
|
|
|
func TestNGramEmbedder_Dim(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := NewNGram(128)
|
|
if got := e.Dim(); got != 128 {
|
|
t.Fatalf("Dim() = %d, want 128", got)
|
|
}
|
|
}
|
|
|
|
func TestNGramEmbedder_DefaultDim(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := NewNGram(0)
|
|
if got := e.Dim(); got != defaultNGramDim {
|
|
t.Fatalf("Dim() = %d, want defaultNGramDim (%d)", got, defaultNGramDim)
|
|
}
|
|
}
|
|
|
|
func TestNGramEmbedder_VectorLengthMatchesDim(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := NewNGram(64)
|
|
vectors, err := e.Embed(context.Background(), []string{"alpha", "beta gamma"})
|
|
if err != nil {
|
|
t.Fatalf("Embed() error: %v", err)
|
|
}
|
|
if len(vectors) != 2 {
|
|
t.Fatalf("got %d vectors, want 2", len(vectors))
|
|
}
|
|
for i, v := range vectors {
|
|
if got := len(v); got != e.Dim() {
|
|
t.Errorf("vector %d length = %d, want %d", i, got, e.Dim())
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestNGramEmbedder_Deterministic(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := NewNGram(128)
|
|
ctx := context.Background()
|
|
texts := []string{"hello world", "the quick brown fox jumps", "Golang embedding"}
|
|
|
|
first, err := e.Embed(ctx, texts)
|
|
if err != nil {
|
|
t.Fatalf("Embed() error: %v", err)
|
|
}
|
|
second, err := e.Embed(ctx, texts)
|
|
if err != nil {
|
|
t.Fatalf("Embed() error: %v", err)
|
|
}
|
|
if len(first) != len(second) {
|
|
t.Fatalf("got %d vectors on first call, %d on second", len(first), len(second))
|
|
}
|
|
for i := range first {
|
|
if !equalVector(first[i], second[i]) {
|
|
t.Fatalf("vector %d differs between calls", i)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestNGramEmbedder_Normalized(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := NewNGram(128)
|
|
vectors, err := e.Embed(context.Background(), []string{
|
|
"a",
|
|
"the quick brown fox jumps over the lazy dog",
|
|
"golang fts embedding test",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Embed() error: %v", err)
|
|
}
|
|
for i, v := range vectors {
|
|
norm := l2Norm(v)
|
|
if math.Abs(norm-1) > 1e-6 {
|
|
t.Errorf("vector %d L2 norm = %f, want 1", i, norm)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestNGramEmbedder_DifferentTextsDiffer(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := NewNGram(128)
|
|
vectors, err := e.Embed(context.Background(), []string{"dog", "cat", "astronaut"})
|
|
if err != nil {
|
|
t.Fatalf("Embed() error: %v", err)
|
|
}
|
|
if len(vectors) != 3 {
|
|
t.Fatalf("got %d vectors, want 3", len(vectors))
|
|
}
|
|
for i := 0; i < len(vectors); i++ {
|
|
if l2Norm(vectors[i]) == 0 {
|
|
t.Errorf("vector %d is a zero vector", i)
|
|
}
|
|
for j := i + 1; j < len(vectors); j++ {
|
|
if equalVector(vectors[i], vectors[j]) {
|
|
t.Errorf("vectors %d and %d are identical", i, j)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestNGramEmbedder_EmptyInput(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := NewNGram(128)
|
|
vectors, err := e.Embed(context.Background(), nil)
|
|
if err != nil {
|
|
t.Fatalf("Embed(nil) error: %v", err)
|
|
}
|
|
if len(vectors) != 0 {
|
|
t.Errorf("got %d vectors, want 0", len(vectors))
|
|
}
|
|
}
|
|
|
|
func TestNGramEmbedder_CanceledContext(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := NewNGram(128)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if _, err := e.Embed(ctx, []string{"hello"}); err == nil {
|
|
t.Fatal("expected error on canceled context, got nil")
|
|
}
|
|
}
|
|
|
|
func equalVector(a, b Vector) bool {
|
|
if len(a) != len(b) {
|
|
return false
|
|
}
|
|
for i := range a {
|
|
if a[i] != b[i] {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func l2Norm(v Vector) float64 {
|
|
var sum float64
|
|
for _, x := range v {
|
|
sum += float64(x) * float64(x)
|
|
}
|
|
return math.Sqrt(sum)
|
|
}
|