118 lines
2.4 KiB
Go
118 lines
2.4 KiB
Go
package utils
|
|
|
|
import (
|
|
"fmt"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
func parseVectorString(v string) ([]float64, error) {
|
|
s := strings.Trim(v, "[]")
|
|
parts := strings.Split(s, ",")
|
|
vec := make([]float64, len(parts))
|
|
for i, p := range parts {
|
|
f, err := strconv.ParseFloat(strings.TrimSpace(p), 64)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
vec[i] = f
|
|
}
|
|
return vec, nil
|
|
}
|
|
|
|
func ParseEmbedding(value any) ([]float64, error) {
|
|
switch v := value.(type) {
|
|
case []float64:
|
|
return v, nil
|
|
|
|
case []any:
|
|
vec := make([]float64, len(v))
|
|
for i, x := range v {
|
|
f, ok := x.(float64)
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid element type %T in embedding", x)
|
|
}
|
|
vec[i] = f
|
|
}
|
|
return vec, nil
|
|
|
|
case string:
|
|
return parseVectorString(v)
|
|
|
|
case []byte:
|
|
return parseVectorString(string(v))
|
|
|
|
default:
|
|
return nil, fmt.Errorf("unexpected embedding type: %T", v)
|
|
}
|
|
}
|
|
|
|
func FormatVector(v []float64) string {
|
|
parts := make([]string, len(v))
|
|
for i, x := range v {
|
|
parts[i] = strconv.FormatFloat(x, 'f', -1, 64)
|
|
}
|
|
return "[" + strings.Join(parts, ",") + "]"
|
|
}
|
|
|
|
// AverageVector returns the element-wise mean of vs. Zero-length vectors (e.g. an
|
|
// embedding model returning an empty vector for a post without text) are skipped, and
|
|
// the accumulator is sized from the first non-empty vector so a leading empty vector
|
|
// cannot panic later in the loop. Returns nil when no usable vector is present.
|
|
func AverageVector(vs [][]float64) []float64 {
|
|
dim := 0
|
|
var valid [][]float64
|
|
for _, v := range vs {
|
|
if len(v) == 0 {
|
|
continue
|
|
}
|
|
if dim == 0 {
|
|
dim = len(v)
|
|
}
|
|
valid = append(valid, v)
|
|
}
|
|
if len(valid) == 0 {
|
|
return nil
|
|
}
|
|
|
|
avg := make([]float64, dim)
|
|
for _, v := range valid {
|
|
for i := 0; i < dim && i < len(v); i++ {
|
|
avg[i] += v[i]
|
|
}
|
|
}
|
|
for i := range avg {
|
|
avg[i] /= float64(len(valid))
|
|
}
|
|
return avg
|
|
}
|
|
|
|
func TopWords(posts []string, stopWords map[string]struct{}, limit int) []string {
|
|
counts := map[string]int{}
|
|
for _, p := range posts {
|
|
for w := range strings.FieldsSeq(strings.ToLower(p)) {
|
|
if len(w) < 4 {
|
|
continue
|
|
}
|
|
if _, blocked := stopWords[w]; blocked {
|
|
continue
|
|
}
|
|
counts[w]++
|
|
}
|
|
}
|
|
type kv struct {
|
|
Word string
|
|
Count int
|
|
}
|
|
arr := make([]kv, 0, len(counts))
|
|
for w, c := range counts {
|
|
arr = append(arr, kv{w, c})
|
|
}
|
|
sort.Slice(arr, func(i, j int) bool { return arr[i].Count > arr[j].Count })
|
|
out := []string{}
|
|
for i := 0; i < len(arr) && i < limit; i++ {
|
|
out = append(out, arr[i].Word)
|
|
}
|
|
return out
|
|
}
|