452 lines
11 KiB
Go
452 lines
11 KiB
Go
package recommendations
|
|
|
|
import (
|
|
"math"
|
|
"math/rand"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// clusterItem represents a single item to be clustered, carrying its embedding,
|
|
// signal weight, and genre metadata.
|
|
type clusterItem struct {
|
|
itemID string
|
|
embedding []float32
|
|
weight float64
|
|
genres []string
|
|
}
|
|
|
|
// kmeansMaxIterations is the maximum number of k-means refinement iterations.
|
|
const kmeansMaxIterations = 50
|
|
|
|
// kmeansConvergenceThreshold is the minimum centroid movement to continue iterating.
|
|
const kmeansConvergenceThreshold = 1e-6
|
|
|
|
// kmeansMinClusterSize is the minimum number of members a cluster must have
|
|
// before it gets merged into a neighbor.
|
|
const kmeansMinClusterSize = 3
|
|
|
|
// kmeansCluster partitions items into k groups using k-means with k-means++
|
|
// initialization. Centroids are computed as weighted averages using each item's
|
|
// signal weight and are L2-normalized after each update. Returns a slice of
|
|
// length len(items) mapping each item index to its assigned cluster index.
|
|
func kmeansCluster(items []clusterItem, k int) []int {
|
|
n := len(items)
|
|
if n == 0 || k <= 0 {
|
|
return nil
|
|
}
|
|
if k > n {
|
|
k = n
|
|
}
|
|
|
|
dims := len(items[0].embedding)
|
|
rng := rand.New(rand.NewSource(kmeansSeed(items, k)))
|
|
|
|
// --- k-means++ initialization ---
|
|
// Select the first centroid uniformly at random.
|
|
centroids := make([][]float64, 0, k)
|
|
firstIdx := rng.Intn(n)
|
|
centroids = append(centroids, float32ToFloat64(items[firstIdx].embedding))
|
|
|
|
// Select remaining centroids with probability proportional to D(x)^2.
|
|
distSq := make([]float64, n)
|
|
for c := 1; c < k; c++ {
|
|
totalDist := 0.0
|
|
for i := range items {
|
|
minD := math.MaxFloat64
|
|
for _, cent := range centroids {
|
|
d := l2DistSq(items[i].embedding, cent)
|
|
if d < minD {
|
|
minD = d
|
|
}
|
|
}
|
|
distSq[i] = minD
|
|
totalDist += minD
|
|
}
|
|
|
|
// Weighted random selection.
|
|
threshold := rng.Float64() * totalDist
|
|
cumulative := 0.0
|
|
chosen := 0
|
|
for i, d := range distSq {
|
|
cumulative += d
|
|
if cumulative >= threshold {
|
|
chosen = i
|
|
break
|
|
}
|
|
}
|
|
centroids = append(centroids, float32ToFloat64(items[chosen].embedding))
|
|
}
|
|
|
|
assignments := make([]int, n)
|
|
|
|
// --- Main k-means loop ---
|
|
for iter := 0; iter < kmeansMaxIterations; iter++ {
|
|
// Assignment step: assign each item to the nearest centroid.
|
|
for i := range items {
|
|
bestCluster := 0
|
|
bestDist := math.MaxFloat64
|
|
for c := range centroids {
|
|
d := l2DistSq(items[i].embedding, centroids[c])
|
|
if d < bestDist {
|
|
bestDist = d
|
|
bestCluster = c
|
|
}
|
|
}
|
|
assignments[i] = bestCluster
|
|
}
|
|
|
|
// Update step: recompute centroids as weighted averages, then L2-normalize.
|
|
newCentroids := make([][]float64, k)
|
|
for c := 0; c < k; c++ {
|
|
newCentroids[c] = make([]float64, dims)
|
|
}
|
|
|
|
for i, ci := range assignments {
|
|
w := items[i].weight
|
|
if w <= 0 {
|
|
w = 0.01 // small floor so every item contributes
|
|
}
|
|
for d := 0; d < dims; d++ {
|
|
newCentroids[ci][d] += float64(items[i].embedding[d]) * w
|
|
}
|
|
}
|
|
|
|
// L2-normalize each centroid.
|
|
for c := 0; c < k; c++ {
|
|
l2NormalizeFloat64(newCentroids[c])
|
|
}
|
|
|
|
// Check convergence: if all centroids moved less than the threshold, stop.
|
|
maxMovement := 0.0
|
|
for c := 0; c < k; c++ {
|
|
movement := 0.0
|
|
for d := 0; d < dims; d++ {
|
|
diff := newCentroids[c][d] - centroids[c][d]
|
|
movement += diff * diff
|
|
}
|
|
movement = math.Sqrt(movement)
|
|
if movement > maxMovement {
|
|
maxMovement = movement
|
|
}
|
|
}
|
|
|
|
centroids = newCentroids
|
|
|
|
if maxMovement < kmeansConvergenceThreshold {
|
|
break
|
|
}
|
|
}
|
|
|
|
// Final assignment pass to ensure consistency with final centroids.
|
|
for i := range items {
|
|
bestCluster := 0
|
|
bestDist := math.MaxFloat64
|
|
for c := range centroids {
|
|
d := l2DistSq(items[i].embedding, centroids[c])
|
|
if d < bestDist {
|
|
bestDist = d
|
|
bestCluster = c
|
|
}
|
|
}
|
|
assignments[i] = bestCluster
|
|
}
|
|
|
|
return assignments
|
|
}
|
|
|
|
func kmeansSeed(items []clusterItem, k int) int64 {
|
|
seed := int64(1469598103934665603)
|
|
seed = seed*1099511628211 + int64(k)
|
|
for _, item := range items {
|
|
for _, b := range []byte(item.itemID) {
|
|
seed ^= int64(b)
|
|
seed *= 1099511628211
|
|
}
|
|
seed ^= int64(math.Float64bits(item.weight))
|
|
seed *= 1099511628211
|
|
}
|
|
if seed == 0 {
|
|
return 1
|
|
}
|
|
return seed
|
|
}
|
|
|
|
// determinClusterCount maps the number of items to a suitable cluster count.
|
|
func determinClusterCount(itemCount int) int {
|
|
switch {
|
|
case itemCount < 10:
|
|
return 1
|
|
case itemCount < 20:
|
|
return 2
|
|
case itemCount <= 30:
|
|
return 3
|
|
case itemCount < 60:
|
|
return 3
|
|
case itemCount <= 100:
|
|
return 4
|
|
case itemCount < 200:
|
|
return 4
|
|
default:
|
|
return 5
|
|
}
|
|
}
|
|
|
|
// mergSmallClusters reassigns items in clusters with fewer than kmeansMinClusterSize
|
|
// members to the nearest non-small cluster. Returns updated assignments with
|
|
// cluster indices compacted to be contiguous starting from 0.
|
|
func mergSmallClusters(items []clusterItem, assignments []int, k int) []int {
|
|
if len(items) == 0 {
|
|
return assignments
|
|
}
|
|
|
|
// Count members per cluster.
|
|
clusterSize := make(map[int]int)
|
|
for _, c := range assignments {
|
|
clusterSize[c]++
|
|
}
|
|
|
|
// Compute centroid of each cluster (weighted).
|
|
dims := len(items[0].embedding)
|
|
centroids := make(map[int][]float64)
|
|
for c := range clusterSize {
|
|
centroids[c] = make([]float64, dims)
|
|
}
|
|
for i, c := range assignments {
|
|
w := items[i].weight
|
|
if w <= 0 {
|
|
w = 0.01
|
|
}
|
|
for d := 0; d < dims; d++ {
|
|
centroids[c][d] += float64(items[i].embedding[d]) * w
|
|
}
|
|
}
|
|
for c := range centroids {
|
|
l2NormalizeFloat64(centroids[c])
|
|
}
|
|
|
|
// Identify small clusters and valid (non-small) clusters.
|
|
smallClusters := make(map[int]bool)
|
|
for c, size := range clusterSize {
|
|
if size < kmeansMinClusterSize {
|
|
smallClusters[c] = true
|
|
}
|
|
}
|
|
|
|
// If all clusters are small, skip merging to avoid losing all data.
|
|
nonSmallCount := 0
|
|
for c := range clusterSize {
|
|
if !smallClusters[c] {
|
|
nonSmallCount++
|
|
}
|
|
}
|
|
if nonSmallCount == 0 {
|
|
return compactAssignments(assignments)
|
|
}
|
|
|
|
// Reassign items from small clusters to the nearest non-small cluster.
|
|
for i, c := range assignments {
|
|
if !smallClusters[c] {
|
|
continue
|
|
}
|
|
bestCluster := -1
|
|
bestDist := math.MaxFloat64
|
|
for target, cent := range centroids {
|
|
if smallClusters[target] {
|
|
continue
|
|
}
|
|
d := l2DistSq(items[i].embedding, cent)
|
|
if d < bestDist {
|
|
bestDist = d
|
|
bestCluster = target
|
|
}
|
|
}
|
|
if bestCluster >= 0 {
|
|
assignments[i] = bestCluster
|
|
}
|
|
}
|
|
|
|
return compactAssignments(assignments)
|
|
}
|
|
|
|
// compactAssignments remaps cluster indices to be contiguous starting from 0.
|
|
func compactAssignments(assignments []int) []int {
|
|
seen := make(map[int]int)
|
|
next := 0
|
|
result := make([]int, len(assignments))
|
|
for i, c := range assignments {
|
|
if _, ok := seen[c]; !ok {
|
|
seen[c] = next
|
|
next++
|
|
}
|
|
result[i] = seen[c]
|
|
}
|
|
return result
|
|
}
|
|
|
|
// buildTasteClusters is the main entry point for computing taste sub-profiles.
|
|
// It determines the appropriate cluster count, runs weighted k-means clustering,
|
|
// merges undersized clusters, and returns labeled TasteCluster values with
|
|
// computed embeddings, dominant genres, and aggregate statistics.
|
|
func buildTasteClusters(items []clusterItem) []TasteCluster {
|
|
if len(items) == 0 {
|
|
return nil
|
|
}
|
|
|
|
// Filter out items without embeddings.
|
|
valid := make([]clusterItem, 0, len(items))
|
|
for _, item := range items {
|
|
if len(item.embedding) > 0 {
|
|
valid = append(valid, item)
|
|
}
|
|
}
|
|
if len(valid) == 0 {
|
|
return nil
|
|
}
|
|
|
|
// Step 1: Determine how many clusters to target.
|
|
k := determinClusterCount(len(valid))
|
|
|
|
// Step 2: Run k-means clustering.
|
|
assignments := kmeansCluster(valid, k)
|
|
|
|
// Step 3: Merge clusters that are too small.
|
|
assignments = mergSmallClusters(valid, assignments, k)
|
|
|
|
// Step 4: Build the output TasteCluster for each cluster.
|
|
clusterMap := make(map[int][]int) // cluster index -> item indices
|
|
for i, c := range assignments {
|
|
clusterMap[c] = append(clusterMap[c], i)
|
|
}
|
|
|
|
// Collect cluster indices and sort for deterministic output order.
|
|
clusterIndices := make([]int, 0, len(clusterMap))
|
|
for c := range clusterMap {
|
|
clusterIndices = append(clusterIndices, c)
|
|
}
|
|
sort.Ints(clusterIndices)
|
|
|
|
clusters := make([]TasteCluster, 0, len(clusterIndices))
|
|
for _, ci := range clusterIndices {
|
|
memberIndices := clusterMap[ci]
|
|
|
|
// Gather embeddings and weights for the weighted average.
|
|
vecs := make([][]float32, len(memberIndices))
|
|
weights := make([]float64, len(memberIndices))
|
|
genreCounts := make(map[string]int)
|
|
totalWeight := 0.0
|
|
|
|
for j, idx := range memberIndices {
|
|
vecs[j] = valid[idx].embedding
|
|
weights[j] = valid[idx].weight
|
|
totalWeight += valid[idx].weight
|
|
for _, g := range valid[idx].genres {
|
|
genreCounts[g]++
|
|
}
|
|
}
|
|
|
|
// Compute L2-normalized weighted average embedding.
|
|
embedding := weightedAverage(vecs, weights)
|
|
|
|
// Determine dominant genres (top 3 by frequency).
|
|
dominantGenres := topNGenres(genreCounts, 3)
|
|
|
|
// Build a human-readable label from the dominant genres.
|
|
label := buildClusterLabel(dominantGenres)
|
|
|
|
clusters = append(clusters, TasteCluster{
|
|
ClusterIdx: ci,
|
|
Embedding: embedding,
|
|
DominantGenres: dominantGenres,
|
|
Label: label,
|
|
MemberCount: len(memberIndices),
|
|
TotalWeight: totalWeight,
|
|
UpdatedAt: time.Now(),
|
|
})
|
|
}
|
|
|
|
return clusters
|
|
}
|
|
|
|
// topNGenres returns up to n genre names sorted by descending frequency.
|
|
func topNGenres(counts map[string]int, n int) []string {
|
|
type genreCount struct {
|
|
genre string
|
|
count int
|
|
}
|
|
|
|
pairs := make([]genreCount, 0, len(counts))
|
|
for g, c := range counts {
|
|
pairs = append(pairs, genreCount{g, c})
|
|
}
|
|
|
|
sort.Slice(pairs, func(i, j int) bool {
|
|
if pairs[i].count != pairs[j].count {
|
|
return pairs[i].count > pairs[j].count
|
|
}
|
|
return pairs[i].genre < pairs[j].genre // stable tie-breaking
|
|
})
|
|
|
|
end := n
|
|
if end > len(pairs) {
|
|
end = len(pairs)
|
|
}
|
|
|
|
result := make([]string, end)
|
|
for i := 0; i < end; i++ {
|
|
result[i] = pairs[i].genre
|
|
}
|
|
return result
|
|
}
|
|
|
|
// buildClusterLabel creates a display label by joining genre names with " & ".
|
|
// Returns "Mixed" if no genres are available.
|
|
func buildClusterLabel(genres []string) string {
|
|
if len(genres) == 0 {
|
|
return "Mixed"
|
|
}
|
|
return strings.Join(genres, " & ")
|
|
}
|
|
|
|
// --- Vector math helpers ---
|
|
|
|
// float32ToFloat64 converts a float32 slice to float64.
|
|
func float32ToFloat64(v []float32) []float64 {
|
|
out := make([]float64, len(v))
|
|
for i, val := range v {
|
|
out[i] = float64(val)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// l2DistSq computes the squared L2 distance between a float32 vector and a
|
|
// float64 centroid. If either vector is nil, returns 0.
|
|
func l2DistSq(a []float32, b []float64) float64 {
|
|
if a == nil || b == nil {
|
|
return 0
|
|
}
|
|
sum := 0.0
|
|
for i := range a {
|
|
diff := float64(a[i]) - b[i]
|
|
sum += diff * diff
|
|
}
|
|
return sum
|
|
}
|
|
|
|
// l2NormalizeFloat64 normalizes a float64 vector in-place to unit length.
|
|
// If the vector has zero magnitude, it is left unchanged.
|
|
func l2NormalizeFloat64(v []float64) {
|
|
var norm float64
|
|
for _, val := range v {
|
|
norm += val * val
|
|
}
|
|
norm = math.Sqrt(norm)
|
|
if norm == 0 {
|
|
return
|
|
}
|
|
for i := range v {
|
|
v[i] /= norm
|
|
}
|
|
}
|