148 lines
4.1 KiB
Go
148 lines
4.1 KiB
Go
package embedding
|
||
|
||
import (
|
||
"context"
|
||
"math"
|
||
"strings"
|
||
)
|
||
|
||
// Default n-gram embedding parameters. These are the zero-configuration values
|
||
// used when NewNGram is called with a non-positive dimension and no options.
|
||
const (
|
||
defaultNGramDim = 256
|
||
defaultNGramMin = 2
|
||
defaultNGramMax = 4
|
||
)
|
||
|
||
// FNV-1a constants (Fowler–Noll–Vo, 64-bit). They give a fixed, cross-process-
|
||
// stable hash so n-gram bucket positions are deterministic across runs and
|
||
// machines (unlike maphash, which randomizes per process).
|
||
const (
|
||
fnv64Offset = 14695981039346656037
|
||
fnv64Prime = 1099511628211
|
||
// signBit is the top bit used to derive the feature-hashing sign.
|
||
signBit = uint64(1) << 63
|
||
)
|
||
|
||
// NGramEmbedder is a zero-dependency feature-hashing embedder. It tokenizes a
|
||
// text into character n-grams, hashes each n-gram into a fixed-size bucket
|
||
// using signed FNV-1a, and L2-normalizes the resulting count vector.
|
||
//
|
||
// It is the default embedder: deterministic, offline, and dependency-free. It
|
||
// is the keyword-independent, semantic-level counterpart to the FTS Index.
|
||
type NGramEmbedder struct {
|
||
dim int
|
||
minN int
|
||
maxN int
|
||
}
|
||
|
||
// NGramOption customizes an NGramEmbedder built by NewNGram.
|
||
type NGramOption func(*NGramEmbedder)
|
||
|
||
// WithNGramRange sets the inclusive character n-gram sizes hashed into the
|
||
// vector. Non-positive or inverted ranges are ignored and the defaults kept.
|
||
func WithNGramRange(minN, maxN int) NGramOption {
|
||
return func(e *NGramEmbedder) {
|
||
if minN > 0 && maxN >= minN {
|
||
e.minN, e.maxN = minN, maxN
|
||
}
|
||
}
|
||
}
|
||
|
||
// NewNGram builds an NGramEmbedder with dimension dim. dim must be positive;
|
||
// non-positive values select defaultNGramDim. The n-gram range defaults to
|
||
// [defaultNGramMin, defaultNGramMax].
|
||
func NewNGram(dim int, opts ...NGramOption) *NGramEmbedder {
|
||
if dim <= 0 {
|
||
dim = defaultNGramDim
|
||
}
|
||
e := &NGramEmbedder{dim: dim, minN: defaultNGramMin, maxN: defaultNGramMax}
|
||
for _, opt := range opts {
|
||
opt(e)
|
||
}
|
||
return e
|
||
}
|
||
|
||
// Dim returns the fixed vector dimension.
|
||
func (e *NGramEmbedder) Dim() int { return e.dim }
|
||
|
||
// Embed returns one L2-normalized n-gram hash vector per input text. The order
|
||
// of the returned vectors matches texts; an empty input yields an empty slice.
|
||
func (e *NGramEmbedder) Embed(ctx context.Context, texts []string) ([]Vector, error) {
|
||
if err := ctx.Err(); err != nil {
|
||
return nil, err
|
||
}
|
||
out := make([]Vector, 0, len(texts))
|
||
for _, text := range texts {
|
||
if err := ctx.Err(); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, e.vectorFor(text))
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// vectorFor builds the n-gram hash vector for one text and L2-normalizes it.
|
||
// Empty text produces a zero vector (left as-is by l2Normalize). Text shorter
|
||
// than the smallest n-gram window hashes the whole text as a single token so it
|
||
// never collapses to a zero vector.
|
||
func (e *NGramEmbedder) vectorFor(text string) Vector {
|
||
v := make(Vector, e.dim)
|
||
lower := strings.ToLower(strings.TrimSpace(text))
|
||
runes := []rune(lower)
|
||
n := len(runes)
|
||
if n == 0 {
|
||
return v
|
||
}
|
||
hashed := false
|
||
for size := e.minN; size <= e.maxN && size <= n; size++ {
|
||
for i := 0; i+size <= n; i++ {
|
||
h := fnv1a64(string(runes[i : i+size]))
|
||
v[int(h%uint64(e.dim))] += signedHash(h)
|
||
hashed = true
|
||
}
|
||
}
|
||
if !hashed {
|
||
h := fnv1a64(lower)
|
||
v[int(h%uint64(e.dim))] += signedHash(h)
|
||
}
|
||
l2Normalize(v)
|
||
return v
|
||
}
|
||
|
||
// fnv1a64 returns the FNV-1a 64-bit hash of s. It is deterministic across
|
||
// processes and platforms, so the same text always lands in the same buckets.
|
||
func fnv1a64(s string) uint64 {
|
||
h := uint64(fnv64Offset)
|
||
for i := 0; i < len(s); i++ {
|
||
h ^= uint64(s[i])
|
||
h *= fnv64Prime
|
||
}
|
||
return h
|
||
}
|
||
|
||
// signedHash maps the top bit of a hash to the +1/-1 sign used to reduce
|
||
// collision bias in feature hashing; the lower bits select the bucket.
|
||
func signedHash(h uint64) float32 {
|
||
if h&signBit != 0 {
|
||
return -1
|
||
}
|
||
return 1
|
||
}
|
||
|
||
// l2Normalize scales v to unit length in place. A zero vector is left as-is
|
||
// rather than dividing by zero.
|
||
func l2Normalize(v Vector) {
|
||
var sum float64
|
||
for _, x := range v {
|
||
sum += float64(x) * float64(x)
|
||
}
|
||
if sum == 0 {
|
||
return
|
||
}
|
||
norm := float32(math.Sqrt(sum))
|
||
for i := range v {
|
||
v[i] /= norm
|
||
}
|
||
}
|