Add timestamps to feedback in Dataset and update related methods and tests

Author: zhenghaoz

dataset/dataset.go

@@ -79,6 +79,7 @@ type Dataset struct {
 	itemLabels   *Labels
 	userFeedback [][]int32
 	itemFeedback [][]int32
+	timestamps   [][]time.Time
 	negatives    [][]int32
 	userDict     *FreqDict
 	itemDict     *FreqDict
@@ -95,6 +96,7 @@ func NewDataset(timestamp time.Time, userCount, itemCount int) *Dataset {
 		itemLabels:   NewLabels(),
 		userFeedback: make([][]int32, userCount),
 		itemFeedback: make([][]int32, itemCount),
+		timestamps:   make([][]time.Time, userCount),
 		userDict:     NewFreqDict(),
 		itemDict:     NewFreqDict(),
 		categories:   make(map[string]int),
@@ -203,6 +205,9 @@ func (d *Dataset) AddUser(user data.User) {
 	if len(d.userFeedback) < len(d.users) {
 		d.userFeedback = append(d.userFeedback, nil)
 	}
+	if len(d.timestamps) < len(d.users) {
+		d.timestamps = append(d.timestamps, nil)
+	}
 }
 
 func (d *Dataset) AddItem(item data.Item) {
@@ -223,11 +228,12 @@ func (d *Dataset) AddItem(item data.Item) {
 	}
 }
 
-func (d *Dataset) AddFeedback(userId, itemId string) {
+func (d *Dataset) AddFeedback(userId, itemId string, timestamp time.Time) {
 	userIndex := d.userDict.Add(userId)
 	itemIndex := d.itemDict.Add(itemId)
 	d.userFeedback[userIndex] = append(d.userFeedback[userIndex], itemIndex)
 	d.itemFeedback[itemIndex] = append(d.itemFeedback[itemIndex], userIndex)
+	d.timestamps[userIndex] = append(d.timestamps[userIndex], timestamp)
 	d.numFeedback++
 }
 
@@ -253,6 +259,7 @@ func (d *Dataset) SplitCF(numTestUsers int, seed int64) (CFSplit, CFSplit) {
 	trainSet.items, testSet.items = d.items, d.items
 	trainSet.userFeedback, testSet.userFeedback = make([][]int32, d.CountUsers()), make([][]int32, d.CountUsers())
 	trainSet.itemFeedback, testSet.itemFeedback = make([][]int32, d.CountItems()), make([][]int32, d.CountItems())
+	trainSet.timestamps, testSet.timestamps = make([][]time.Time, d.CountUsers()), make([][]time.Time, d.CountUsers())
 	trainSet.userDict, testSet.userDict = d.userDict, d.userDict
 	trainSet.itemDict, testSet.itemDict = d.itemDict, d.itemDict
 	rng := util.NewRandomGenerator(seed)
@@ -262,11 +269,13 @@ func (d *Dataset) SplitCF(numTestUsers int, seed int64) (CFSplit, CFSplit) {
 				k := rng.Intn(len(d.userFeedback[userIndex]))
 				testSet.userFeedback[userIndex] = append(testSet.userFeedback[userIndex], d.userFeedback[userIndex][k])
 				testSet.itemFeedback[d.userFeedback[userIndex][k]] = append(testSet.itemFeedback[d.userFeedback[userIndex][k]], userIndex)
+				testSet.timestamps[userIndex] = append(testSet.timestamps[userIndex], d.timestamps[userIndex][k])
 				testSet.numFeedback++
 				for i, itemIndex := range d.userFeedback[userIndex] {
 					if i != k {
 						trainSet.userFeedback[userIndex] = append(trainSet.userFeedback[userIndex], itemIndex)
 						trainSet.itemFeedback[itemIndex] = append(trainSet.itemFeedback[itemIndex], userIndex)
+						trainSet.timestamps[userIndex] = append(trainSet.timestamps[userIndex], d.timestamps[userIndex][i])
 						trainSet.numFeedback++
 					}
 				}
@@ -277,13 +286,15 @@ func (d *Dataset) SplitCF(numTestUsers int, seed int64) (CFSplit, CFSplit) {
 		for _, userIndex := range testUsers {
 			if len(d.userFeedback[userIndex]) > 0 {
 				k := rng.Intn(len(d.userFeedback[userIndex]))
-				testSet.userFeedback[userIndex] = append(testSet.userFeedback[userIndex], d.userFeedback[userIndex][k])
+				testSet.timestamps[userIndex] = append(testSet.timestamps[userIndex], d.timestamps[userIndex][k])
 				testSet.itemFeedback[d.userFeedback[userIndex][k]] = append(testSet.itemFeedback[d.userFeedback[userIndex][k]], userIndex)
+				testSet.userFeedback[userIndex] = append(testSet.userFeedback[userIndex], d.userFeedback[userIndex][k])
 				testSet.numFeedback++
 				for i, itemIndex := range d.userFeedback[userIndex] {
 					if i != k {
 						trainSet.userFeedback[userIndex] = append(trainSet.userFeedback[userIndex], itemIndex)
 						trainSet.itemFeedback[itemIndex] = append(trainSet.itemFeedback[itemIndex], userIndex)
+						trainSet.timestamps[userIndex] = append(trainSet.timestamps[userIndex], d.timestamps[userIndex][i])
 						trainSet.numFeedback++
 					}
 				}
@@ -292,9 +303,10 @@ func (d *Dataset) SplitCF(numTestUsers int, seed int64) (CFSplit, CFSplit) {
 		testUserSet := mapset.NewSet(testUsers...)
 		for userIndex := int32(0); userIndex < int32(d.CountUsers()); userIndex++ {
 			if !testUserSet.Contains(userIndex) {
-				for _, itemIndex := range d.userFeedback[userIndex] {
+				for idx, itemIndex := range d.userFeedback[userIndex] {
 					trainSet.userFeedback[userIndex] = append(trainSet.userFeedback[userIndex], itemIndex)
 					trainSet.itemFeedback[itemIndex] = append(trainSet.itemFeedback[itemIndex], userIndex)
+					trainSet.timestamps[userIndex] = append(trainSet.timestamps[userIndex], d.timestamps[userIndex][idx])
 					trainSet.numFeedback++
 				}
 			}
@@ -303,6 +315,39 @@ func (d *Dataset) SplitCF(numTestUsers int, seed int64) (CFSplit, CFSplit) {
 	return trainSet, testSet
 }
 
+// SplitLatest splits dataset by moving the most recent feedback of all users into the test set to avoid leakage.
+func (d *Dataset) SplitLatest() (CFSplit, CFSplit) {
+	trainSet, testSet := new(Dataset), new(Dataset)
+	trainSet.users, testSet.users = d.users, d.users
+	trainSet.items, testSet.items = d.items, d.items
+	trainSet.userFeedback, testSet.userFeedback = make([][]int32, d.CountUsers()), make([][]int32, d.CountUsers())
+	trainSet.itemFeedback, testSet.itemFeedback = make([][]int32, d.CountItems()), make([][]int32, d.CountItems())
+	trainSet.timestamps, testSet.timestamps = make([][]time.Time, d.CountUsers()), make([][]time.Time, d.CountUsers())
+	trainSet.userDict, testSet.userDict = d.userDict, d.userDict
+	trainSet.itemDict, testSet.itemDict = d.itemDict, d.itemDict
+	for userIndex := int32(0); userIndex < int32(d.CountUsers()); userIndex++ {
+		if len(d.userFeedback[userIndex]) == 0 {
+			continue
+		}
+		_, latestIdx := lo.MaxIndexBy(d.timestamps[userIndex], func(a, b time.Time) bool {
+			return a.After(b)
+		})
+		testSet.timestamps[userIndex] = append(testSet.timestamps[userIndex], d.timestamps[userIndex][latestIdx])
+		testSet.itemFeedback[d.userFeedback[userIndex][latestIdx]] = append(testSet.itemFeedback[d.userFeedback[userIndex][latestIdx]], userIndex)
+		testSet.userFeedback[userIndex] = append(testSet.userFeedback[userIndex], d.userFeedback[userIndex][latestIdx])
+		testSet.numFeedback++
+		for i, itemIndex := range d.userFeedback[userIndex] {
+			if i != latestIdx {
+				trainSet.userFeedback[userIndex] = append(trainSet.userFeedback[userIndex], itemIndex)
+				trainSet.itemFeedback[itemIndex] = append(trainSet.itemFeedback[itemIndex], userIndex)
+				trainSet.timestamps[userIndex] = append(trainSet.timestamps[userIndex], d.timestamps[userIndex][i])
+				trainSet.numFeedback++
+			}
+		}
+	}
+	return trainSet, testSet
+}
+
 type Labels struct {
 	fields *strutil.Pool
 	values *FreqDict
@@ -366,6 +411,7 @@ func LoadDataFromBuiltIn(dataSetName string) (*Dataset, *Dataset, error) {
 	test.userDict, test.itemDict = train.userDict, train.itemDict
 	test.userFeedback = make([][]int32, len(train.userFeedback))
 	test.itemFeedback = make([][]int32, len(train.itemFeedback))
+	test.timestamps = make([][]time.Time, len(train.userFeedback))
 	test.negatives = make([][]int32, len(train.userFeedback))
 	err = loadTest(test, testFilePath)
 	if err != nil {
@@ -404,7 +450,7 @@ func loadTrain(path string) (*Dataset, error) {
 			dataset.AddItem(data.Item{ItemId: util.FormatInt(i)})
 		}
 		// add feedback
-		dataset.AddFeedback(fields[0], fields[1])
+		dataset.AddFeedback(fields[0], fields[1], time.Time{})
 	}
 	return dataset, scanner.Err()
 }
@@ -429,7 +475,7 @@ func loadTest(dataset *Dataset, path string) error {
 		positive = positive[1 : len(positive)-1]
 		fields = strings.Split(positive, ",")
 		// add feedback
-		dataset.AddFeedback(fields[0], fields[1])
+		dataset.AddFeedback(fields[0], fields[1], time.Time{})
 		// add negatives
 		userId, err := strconv.Atoi(fields[0])
 		if err != nil {

dataset/dataset_test.go

@@ -159,14 +159,15 @@ func TestDataset_AddFeedback(t *testing.T) {
 	}
 	for i := 0; i < 10; i++ {
 		for j := i; j < 10; j++ {
-			dataSet.AddFeedback(strconv.Itoa(i), strconv.Itoa(j))
+			dataSet.AddFeedback(strconv.Itoa(i), strconv.Itoa(j), time.Unix(int64(i*10+j), 0))
 		}
 	}
 	userIDF := dataSet.GetUserIDF()
 	itemIDF := dataSet.GetItemIDF()
 	for i := 0; i < 10; i++ {
 		assert.Len(t, dataSet.GetUserFeedback()[i], 10-i)
 		assert.Len(t, dataSet.GetItemFeedback()[i], i+1)
+		assert.Len(t, dataSet.timestamps[i], 10-i)
 		assert.InDelta(t, math32.Log(float32(10)/float32(10-i)), userIDF[i], 1e-2)
 		assert.InDelta(t, math32.Log(float32(10)/float32(i+1)), itemIDF[i], 1e-2)
 	}
@@ -184,7 +185,7 @@ func TestDataset_Split(t *testing.T) {
 	}
 	for i := 0; i < numUsers; i++ {
 		for j := i + 1; j < numItems; j++ {
-			dataset.AddFeedback(fmt.Sprintf("user%v", i), fmt.Sprintf("item%v", j))
+			dataset.AddFeedback(fmt.Sprintf("user%v", i), fmt.Sprintf("item%v", j), time.Time{})
 		}
 	}
 	assert.Equal(t, 9, dataset.CountFeedback())
@@ -206,6 +207,37 @@ func TestDataset_Split(t *testing.T) {
 	assert.Equal(t, 2, test2.CountFeedback())
 }
 
+func TestDataset_SplitLatest(t *testing.T) {
+	const numUsers, numItems = 3, 5
+	// create dataset
+	dataset := NewDataset(time.Now(), numUsers, numItems)
+	for i := 0; i < numUsers; i++ {
+		dataset.AddUser(data.User{UserId: fmt.Sprintf("user%v", i)})
+	}
+	for i := 0; i < numItems; i++ {
+		dataset.AddItem(data.Item{ItemId: fmt.Sprintf("item%v", i)})
+	}
+	for i := 0; i < numUsers; i++ {
+		for j := i + 1; j < numItems; j++ {
+			dataset.AddFeedback(fmt.Sprintf("user%v", i), fmt.Sprintf("item%v", j), time.Unix(int64(j), 0))
+		}
+	}
+	assert.Equal(t, 9, dataset.CountFeedback())
+	// split
+	train, test := dataset.SplitLatest()
+	assert.Equal(t, numUsers, train.CountUsers())
+	assert.Equal(t, numItems, train.CountItems())
+	assert.Equal(t, numUsers, test.CountUsers())
+	assert.Equal(t, numItems, test.CountItems())
+	assert.Equal(t, 6, train.CountFeedback())
+	assert.Equal(t, 3, test.CountFeedback())
+	for i := 0; i < numUsers; i++ {
+		assert.Len(t, train.GetUserFeedback()[i], numItems-i-2)
+		assert.Len(t, test.GetUserFeedback()[i], 1)
+		assert.Equal(t, 4, int(test.GetUserFeedback()[i][0]))
+	}
+}
+
 func TestDataset_LoadMovieLens1M(t *testing.T) {
 	train, test, err := LoadDataFromBuiltIn("ml-1m")
 	assert.NoError(t, err)

master/tasks.go

@@ -479,7 +479,7 @@ func (m *Master) LoadDataFromDatabase(
 						break
 					}
 				}
-				dataSet.AddFeedback(f.UserId, f.ItemId)
+				dataSet.AddFeedback(f.UserId, f.ItemId, f.Timestamp)
 			}
 			span.Add(len(feedback))
 		}

model/cf/evaluator_test.go

@@ -144,7 +144,7 @@ func TestEvaluate(t *testing.T) {
 	}
 	for i := 0; i < 16; i++ {
 		test.AddItem(data.Item{ItemId: strconv.Itoa(i)})
-		test.AddFeedback(strconv.Itoa(i/4), strconv.Itoa(i))
+		test.AddFeedback(strconv.Itoa(i/4), strconv.Itoa(i), time.Time{})
 	}
 	assert.Equal(t, 16, test.CountFeedback())
 	assert.Equal(t, 4, test.CountUsers())