master: unify datasets between tasks

dataset/dataset.go

@@ -15,10 +15,19 @@
 package dataset
 
 import (
+	"bufio"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/chewxy/math32"
+	mapset "github.com/deckarep/golang-set/v2"
+	"github.com/juju/errors"
 	"github.com/samber/lo"
+	"github.com/zhenghaoz/gorse/base"
+	"github.com/zhenghaoz/gorse/model"
 	"github.com/zhenghaoz/gorse/storage/data"
 	"modernc.org/strutil"
 )
@@ -31,10 +40,12 @@ type Dataset struct {
 	items        []data.Item
 	userLabels   *Labels
 	itemLabels   *Labels
-	userFeedback [][]ID
-	itemFeedback [][]ID
+	userFeedback [][]int32
+	itemFeedback [][]int32
+	negatives    [][]int32
 	userDict     *FreqDict
 	itemDict     *FreqDict
+	categories   mapset.Set[string]
 }
 
 func NewDataset(timestamp time.Time, userCount, itemCount int) *Dataset {
@@ -44,10 +55,11 @@ func NewDataset(timestamp time.Time, userCount, itemCount int) *Dataset {
 		items:        make([]data.Item, 0, itemCount),
 		userLabels:   NewLabels(),
 		itemLabels:   NewLabels(),
-		userFeedback: make([][]ID, userCount),
-		itemFeedback: make([][]ID, itemCount),
+		userFeedback: make([][]int32, userCount),
+		itemFeedback: make([][]int32, itemCount),
 		userDict:     NewFreqDict(),
 		itemDict:     NewFreqDict(),
+		categories:   mapset.NewSet[string](),
 	}
 }
 
@@ -59,18 +71,30 @@ func (d *Dataset) GetUsers() []data.User {
 	return d.users
 }
 
+func (d *Dataset) CountUsers() int {
+	return len(d.users)
+}
+
 func (d *Dataset) GetItems() []data.Item {
 	return d.items
 }
 
-func (d *Dataset) GetUserFeedback() [][]ID {
+func (d *Dataset) CountItems() int {
+	return len(d.items)
+}
+
+func (d *Dataset) GetUserFeedback() [][]int32 {
 	return d.userFeedback
 }
 
-func (d *Dataset) GetItemFeedback() [][]ID {
+func (d *Dataset) GetItemFeedback() [][]int32 {
 	return d.itemFeedback
 }
 
+func (d *Dataset) GetCategories() []string {
+	return d.categories.ToSlice()
+}
+
 // GetUserIDF returns the IDF of users.
 //
 //	IDF(u) = log(I/freq(u))
@@ -145,13 +169,27 @@ func (d *Dataset) AddItem(item data.Item) {
 	if len(d.itemFeedback) < len(d.items) {
 		d.itemFeedback = append(d.itemFeedback, nil)
 	}
+	d.categories.Append(item.Categories...)
 }
 
 func (d *Dataset) AddFeedback(userId, itemId string) {
 	userIndex := d.userDict.Id(userId)
 	itemIndex := d.itemDict.Id(itemId)
-	d.userFeedback[userIndex] = append(d.userFeedback[userIndex], ID(itemIndex))
-	d.itemFeedback[itemIndex] = append(d.itemFeedback[itemIndex], ID(userIndex))
+	d.userFeedback[userIndex] = append(d.userFeedback[userIndex], int32(itemIndex))
+	d.itemFeedback[itemIndex] = append(d.itemFeedback[itemIndex], int32(userIndex))
+}
+
+func (d *Dataset) NegativeSample(excludeSet *Dataset, numCandidates int) [][]int32 {
+	if len(d.negatives) == 0 {
+		rng := base.NewRandomGenerator(0)
+		d.negatives = make([][]int32, d.CountUsers())
+		for userIndex := 0; userIndex < d.CountUsers(); userIndex++ {
+			s1 := mapset.NewSet(d.GetUserFeedback()[userIndex]...)
+			s2 := mapset.NewSet(excludeSet.GetUserFeedback()[userIndex]...)
+			d.negatives[userIndex] = rng.SampleInt32(0, int32(d.CountItems()), numCandidates, s1, s2)
+		}
+	}
+	return d.negatives
 }
 
 type Labels struct {
@@ -200,3 +238,96 @@ func isSliceOf[T any](v []any) bool {
 	}
 	return true
 }
+
+func LoadDataFromBuiltIn(dataSetName string) (*Dataset, *Dataset, error) {
+	// Extract Data set information
+	trainFilePath, testFilePath, err := model.LocateBuiltInDataset(dataSetName, model.FormatNCF)
+	if err != nil {
+		return nil, nil, err
+	}
+	// Load dataset
+	train, err := loadTrain(trainFilePath)
+	if err != nil {
+		return nil, nil, err
+	}
+	test := NewDataset(train.GetTimestamp(), 0, 0)
+	test.users, test.items = train.users, train.items
+	test.userDict, test.itemDict = train.userDict, train.itemDict
+	test.userFeedback = make([][]int32, len(train.userFeedback))
+	test.itemFeedback = make([][]int32, len(train.itemFeedback))
+	test.negatives = make([][]int32, len(train.userFeedback))
+	err = loadTest(test, testFilePath)
+	if err != nil {
+		return nil, nil, err
+	}
+	return train, test, nil
+}
+
+func loadTrain(path string) (*Dataset, error) {
+	dataset := NewDataset(time.Now(), 0, 0)
+	// Open
+	file, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+	// Read lines
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := scanner.Text()
+		fields := strings.Split(line, "\t")
+		// add users
+		userId, err := strconv.Atoi(fields[0])
+		if err != nil {
+			return nil, err
+		}
+		for i := dataset.userDict.Count(); i <= userId; i++ {
+			dataset.AddUser(data.User{UserId: strconv.Itoa(i)})
+		}
+		// add items
+		itemId, err := strconv.Atoi(fields[1])
+		if err != nil {
+			return nil, err
+		}
+		for i := dataset.itemDict.Count(); i <= itemId; i++ {
+			dataset.AddItem(data.Item{ItemId: strconv.Itoa(i)})
+		}
+		// add feedback
+		dataset.AddFeedback(fields[0], fields[1])
+	}
+	return dataset, scanner.Err()
+}
+
+func loadTest(dataset *Dataset, path string) error {
+	// Open
+	file, err := os.Open(path)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	defer file.Close()
+	// Read lines
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := scanner.Text()
+		// parse line
+		fields := strings.Split(line, "\t")
+		positive, negatives := fields[0], fields[1:]
+		if positive[0] != '(' || positive[len(positive)-1] != ')' {
+			return fmt.Errorf("wrong foramt: %v", line)
+		}
+		positive = positive[1 : len(positive)-1]
+		fields = strings.Split(positive, ",")
+		// add feedback
+		dataset.AddFeedback(fields[0], fields[1])
+		// add negatives
+		userId, err := strconv.Atoi(fields[0])
+		if err != nil {
+			return err
+		}
+		dataset.negatives[userId] = make([]int32, len(negatives))
+		for i, negative := range negatives {
+			dataset.negatives[userId][i] = int32(dataset.itemDict.Id(negative))
+		}
+	}
+	return scanner.Err()
+}

dataset/dict.go

@@ -42,6 +42,13 @@ func (d *FreqDict) Id(s string) (y int) {
 	return
 }
 
+func (d *FreqDict) NotAdd(s string) int {
+	if y, ok := d.si[s]; ok {
+		return y
+	}
+	return -1
+}
+
 func (d *FreqDict) NotCount(s string) (y int) {
 	if y, ok := d.si[s]; ok {
 		return y

dataset/dict_test.go

@@ -32,4 +32,6 @@ func TestFreqDict(t *testing.T) {
 	assert.Equal(t, 1, dict.Freq(0))
 	assert.Equal(t, 2, dict.Freq(1))
 	assert.Equal(t, 3, dict.Freq(2))
+	assert.Equal(t, 0, dict.NotAdd("a"))
+	assert.Equal(t, -1, dict.NotAdd("e"))
 }

master/tasks.go

@@ -171,7 +171,7 @@
 	InactiveItemsTotal.Set(float64(inactiveItems))
 
 	// write categories to cache
-	if err = m.CacheClient.SetSet(ctx, cache.ItemCategories, rankingDataset.CategorySet.ToSlice()...); err != nil {
+	if err = m.CacheClient.SetSet(ctx, cache.ItemCategories, dataSet.GetCategories()...); err != nil {
 		log.Logger().Error("failed to write categories to cache", zap.Error(err))
 	}
 
@@ -669,10 +669,6 @@
 		temp := time.Now().AddDate(0, 0, -int(positiveFeedbackTTL))
 		feedbackTimeLimit = data.WithBeginTime(temp)
 	}
-	timeWindowLimit := time.Time{}
-	if m.Config.Recommend.Popular.PopularWindow > 0 {
-		timeWindowLimit = time.Now().Add(-m.Config.Recommend.Popular.PopularWindow)
-	}
 	rankingDataset = ranking.NewMapIndexDataset()
 
 	// STEP 1: pull users
@@ -742,9 +738,6 @@
 			itemIndex := rankingDataset.ItemIndex.ToNumber(item.ItemId)
 			if len(rankingDataset.ItemFeatures) == int(itemIndex) {
 				rankingDataset.ItemFeatures = append(rankingDataset.ItemFeatures, nil)
-				rankingDataset.HiddenItems = append(rankingDataset.HiddenItems, false)
-				rankingDataset.ItemCategories = append(rankingDataset.ItemCategories, item.Categories)
-				rankingDataset.CategorySet.Append(item.Categories...)
 			}
 			features := click.ConvertLabelsToFeatures(item.Labels)
 			rankingDataset.NumItemLabelUsed += len(features)
@@ -772,9 +765,6 @@
 					})
 				}
 			}
-			if item.IsHidden { // set hidden flag
-				rankingDataset.HiddenItems[itemIndex] = true
-			}
 			dataSet.AddItem(item)
 		}
 		span.Add(len(batchItems))
@@ -790,7 +780,6 @@
 	LoadDatasetStepSecondsVec.WithLabelValues("load_items").Set(time.Since(start).Seconds())
 
 	// create positive set
-	popularCount := make([]int32, rankingDataset.ItemCount())
 	positiveSet := make([]mapset.Set[int32], rankingDataset.UserCount())
 	for i := range positiveSet {
 		positiveSet[i] = mapset.NewSet[int32]()
@@ -835,10 +824,6 @@
 				posFeedbackCount++
 				// insert feedback to ranking dataset
 				rankingDataset.AddFeedback(f.UserId, f.ItemId, false)
-				// insert feedback to popularity counter
-				if f.Timestamp.After(timeWindowLimit) && !rankingDataset.HiddenItems[itemIndex] {
-					popularCount[itemIndex]++
-				}
 				// insert feedback to evaluator
 				evaluator.Positive(f.FeedbackType, userIndex, itemIndex, f.Timestamp)
 				mu.Unlock()
@@ -1030,7 +1015,7 @@
 				span.Add(1)
 			}
 		}
 	}
 
 	// Save item-to-item recommendations to cache
 	for i, recommender := range itemToItemRecommenders {
@@ -1151,7 +1136,7 @@
 				log.Logger().Error("failed to save user neighbors to cache", zap.String("user_id", user.UserId), zap.Error(err))
 				continue
 			}
 			// Save user-to-user digest and last update time to cache
 			if err := m.CacheClient.Set(ctx,
 				cache.String(cache.Key(cache.UserToUserDigest, cache.Key(cache.Neighbors, user.UserId)), m.Config.UserNeighborDigest()),
 				cache.Time(cache.Key(cache.UserToUserUpdateTime, cache.Key(cache.Neighbors, user.UserId)), time.Now()),

model/ranking/data.go

@@ -31,18 +31,15 @@ import (
 
 // DataSet contains preprocessed data structures for recommendation models.
 type DataSet struct {
-	UserIndex      base.Index
-	ItemIndex      base.Index
-	FeedbackUsers  base.Array[int32]
-	FeedbackItems  base.Array[int32]
-	UserFeedback   [][]int32
-	ItemFeedback   [][]int32
-	Negatives      [][]int32
-	ItemFeatures   [][]lo.Tuple2[int32, float32]
-	UserFeatures   [][]lo.Tuple2[int32, float32]
-	HiddenItems    []bool
-	ItemCategories [][]string
-	CategorySet    mapset.Set[string]
+	UserIndex     base.Index
+	ItemIndex     base.Index
+	FeedbackUsers base.Array[int32]
+	FeedbackItems base.Array[int32]
+	UserFeedback  [][]int32
+	ItemFeedback  [][]int32
+	Negatives     [][]int32
+	ItemFeatures  [][]lo.Tuple2[int32, float32]
+	UserFeatures  [][]lo.Tuple2[int32, float32]
 	// statistics
 	NumItemLabels    int32
 	NumUserLabels    int32
@@ -53,7 +50,6 @@ type DataSet struct {
 // NewMapIndexDataset creates a data set.
 func NewMapIndexDataset() *DataSet {
 	s := new(DataSet)
-	s.CategorySet = mapset.NewSet[string]()
 	// Create index
 	s.UserIndex = base.NewMapIndex()
 	s.ItemIndex = base.NewMapIndex()
@@ -179,9 +175,6 @@ func (dataset *DataSet) Split(numTestUsers int, seed int64) (*DataSet, *DataSet)
 	trainSet, testSet := new(DataSet), new(DataSet)
 	trainSet.NumItemLabels, testSet.NumItemLabels = dataset.NumItemLabels, dataset.NumItemLabels
 	trainSet.NumUserLabels, testSet.NumUserLabels = dataset.NumUserLabels, dataset.NumUserLabels
-	trainSet.HiddenItems, testSet.HiddenItems = dataset.HiddenItems, dataset.HiddenItems
-	trainSet.ItemCategories, testSet.ItemCategories = dataset.ItemCategories, dataset.ItemCategories
-	trainSet.CategorySet, testSet.CategorySet = dataset.CategorySet, dataset.CategorySet
 	trainSet.ItemFeatures, testSet.ItemFeatures = dataset.ItemFeatures, dataset.ItemFeatures
 	trainSet.UserFeatures, testSet.UserFeatures = dataset.UserFeatures, dataset.UserFeatures
 	trainSet.NumItemLabelUsed, testSet.NumItemLabelUsed = dataset.NumItemLabelUsed, dataset.NumItemLabelUsed