disco

package module
v0.2.0 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: May 20, 2025 License: MIT Imports: 16 Imported by: 2

README

Disco Go

🔥 Recommendations for Go using collaborative filtering

  • Supports user-based and item-based recommendations
  • Works with explicit and implicit feedback
  • Uses high-performance matrix factorization

Build Status

Installation

Run:

go get github.com/ankane/disco-go

Getting Started

Import the package

import "github.com/ankane/disco-go"

Prep your data in the format userId, itemId, value

data := disco.NewDataset[string, string]()
data.Push("user_a", "item_a", 5.0)
data.Push("user_a", "item_b", 3.5)
data.Push("user_b", "item_a", 4.0)

IDs can be integers or strings

data.Push(1, "item_a", 5.0)

If users rate items directly, this is known as explicit feedback. Fit the recommender with:

recommender, err := disco.FitExplicit(data)

If users don’t rate items directly (for instance, they’re purchasing items or reading posts), this is known as implicit feedback. Use 1.0 or a value like number of purchases or page views for the dataset, and fit the recommender with:

recommender, err := disco.FitImplicit(data)

Get user-based recommendations - “users like you also liked”

recommender.UserRecs(userId, 5)

Get item-based recommendations - “users who liked this item also liked”

recommender.ItemRecs(itemId, 5)

Get the predicted rating for a specific user and item

recommender.Predict(userId, itemId)

Get similar users

recommender.SimilarUsers(userId, 5)

Examples

MovieLens

Load the data

data, err := disco.LoadMovieLens()

Create a recommender

recommender, err := disco.FitExplicit(data, disco.Factors(20))

Get similar movies

recommender.ItemRecs("Star Wars (1977)", 5)

Storing Recommendations

Save recommendations to your database.

Alternatively, you can store only the factors and use a library like pgvector-go. See an example.

Algorithms

Disco uses high-performance matrix factorization.

Specify the number of factors and iterations

recommender, err := disco.FitExplicit(data, disco.Factors(8), disco.Iterations(20))

Progress

Pass a callback to show progress

callback := func(info disco.FitInfo) { fmt.Printf("%+v\n", info) }
recommender, err := disco.FitExplicit(data, disco.Callback(callback))

Note: TrainLoss and ValidLoss are not available for implicit feedback

Validation

Pass a validation set with explicit feedback

recommender, err := disco.FitEvalExplicit(trainSet, validSet)

The loss function is RMSE

Cold Start

Collaborative filtering suffers from the cold start problem. It’s unable to make good recommendations without data on a user or item, which is problematic for new users and items.

recommender.UserRecs(newUserId, 5) // returns empty array

There are a number of ways to deal with this, but here are some common ones:

  • For user-based recommendations, show new users the most popular items
  • For item-based recommendations, make content-based recommendations

Reference

Get ids

recommender.UserIds()
recommender.ItemIds()

Get the global mean

recommender.GlobalMean()

Get factors

recommender.UserFactors(userId)
recommender.ItemFactors(itemId)

References

History

View the changelog

Contributing

Everyone is encouraged to help improve this project. Here are a few ways you can help:

To get started with development:

git clone https://github.com/ankane/disco-go.git
cd disco-go
go mod tidy
go test -v

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

This section is empty.

Types

type Dataset

type Dataset[T Id, U Id] struct {
	// contains filtered or unexported fields
}

A dataset.

func LoadMovieLens

func LoadMovieLens() (*Dataset[int, string], error)

Loads the MovieLens 100K dataset.

func NewDataset

func NewDataset[T Id, U Id]() *Dataset[T, U]

Creates a new dataset.

func (*Dataset[T, U]) Grow added in v0.2.0

func (d *Dataset[T, U]) Grow(capacity int)

Grows the capacity of the dataset.

func (*Dataset[T, U]) Len

func (d *Dataset[T, U]) Len() int

Returns the number of ratings in the dataset.

func (*Dataset[T, U]) Push

func (d *Dataset[T, U]) Push(userId T, itemId U, value float32)

Adds a rating to the dataset.

func (*Dataset[T, U]) SplitRandom

func (d *Dataset[T, U]) SplitRandom(p float32) (*Dataset[T, U], *Dataset[T, U])

Splits the dataset into training and validation sets.

type FitInfo

type FitInfo struct {
	// The iteration.
	Iteration int
	// The training loss.
	TrainLoss float32
	// The validation loss.
	ValidLoss float32
}

Information about a training iteration.

type Id

type Id interface {
	string | int | uint | int8 | int16 | int32 | int64 | uint8 | uint16 | uint32 | uint64
}

An id.

type Option

type Option func(*config)

A recommender option.

func Alpha added in v0.1.1

func Alpha(alpha float32) Option

Sets alpha.

func Callback

func Callback(callback func(info FitInfo)) Option

Sets the callback for each iteration.

func Factors

func Factors(factors int) Option

Sets the number of factors.

func Iterations

func Iterations(iterations int) Option

Sets the number of iterations.

func LearningRate

func LearningRate(learningRate float32) Option

Sets the learning rate.

func Regularization

func Regularization(regularization float32) Option

Sets the regularization.

func Seed

func Seed(seed int64) Option

Sets the random seed.

type Rec

type Rec[T Id] struct {
	Id    T
	Score float32
}

A recommendation.

type Recommender

type Recommender[T Id, U Id] struct {
	// contains filtered or unexported fields
}

A recommender.

func FitEvalExplicit

func FitEvalExplicit[T Id, U Id](trainSet *Dataset[T, U], validSet *Dataset[T, U], options ...Option) (*Recommender[T, U], error)

Creates a recommender with explicit feedback and performs cross-validation.

func FitExplicit

func FitExplicit[T Id, U Id](trainSet *Dataset[T, U], options ...Option) (*Recommender[T, U], error)

Creates a recommender with explicit feedback.

func FitImplicit

func FitImplicit[T Id, U Id](trainSet *Dataset[T, U], options ...Option) (*Recommender[T, U], error)

Creates a recommender with implicit feedback.

func (*Recommender[T, U]) GlobalMean

func (r *Recommender[T, U]) GlobalMean() float32

Returns the global mean.

func (*Recommender[T, U]) ItemFactors

func (r *Recommender[T, U]) ItemFactors(itemId U) []float32

Returns factors for a specific item.

func (*Recommender[T, U]) ItemIds

func (r *Recommender[T, U]) ItemIds() []U

Returns item ids.

func (*Recommender[T, U]) ItemRecs

func (r *Recommender[T, U]) ItemRecs(itemId U, count int) []Rec[U]

Returns recommendations for an item.

func (*Recommender[T, U]) Predict

func (r *Recommender[T, U]) Predict(userId T, itemId U) float32

Returns the predicted rating for a specific user and item.

func (*Recommender[T, U]) Rmse

func (r *Recommender[T, U]) Rmse(data *Dataset[T, U]) float32

Calculates the root mean square error for a dataset.

func (*Recommender[T, U]) SimilarUsers

func (r *Recommender[T, U]) SimilarUsers(userId T, count int) []Rec[T]

Returns similar users.

func (*Recommender[T, U]) UserFactors

func (r *Recommender[T, U]) UserFactors(userId T) []float32

Returns factors for a specific user.

func (*Recommender[T, U]) UserIds

func (r *Recommender[T, U]) UserIds() []T

Returns user ids.

func (*Recommender[T, U]) UserRecs

func (r *Recommender[T, U]) UserRecs(userId T, count int) []Rec[U]

Returns recommendations for a user.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL