from fastai.collab import * # Quick access to collab filtering functionality
collab models use data in a DataFrame of user, items, and ratings.
path = untar_data(URLs.ML_SAMPLE)
path
PosixPath('/home/ubuntu/.fastai/data/movie_lens_sample')
ratings = pd.read_csv(path/'ratings.csv')
series2cat(ratings, 'userId', 'movieId')
ratings.head()
| userId | movieId | rating | timestamp | |
|---|---|---|---|---|
| 0 | 73 | 1097 | 4.0 | 1255504951 |
| 1 | 561 | 924 | 3.5 | 1172695223 |
| 2 | 157 | 260 | 3.5 | 1291598691 |
| 3 | 358 | 1210 | 5.0 | 957481884 |
| 4 | 130 | 316 | 2.0 | 1138999234 |
data = CollabDataBunch.from_df(ratings, seed=42)
y_range = [0, 5.5]
That's all we need to create and train a model:
learn = collab_learner(data, n_factors=50, y_range=y_range)
learn.fit_one_cycle(4, 5e-3)
| epoch | train_loss | valid_loss | time |
|---|---|---|---|
| 0 | 1.779086 | 1.265616 | 00:00 |
| 1 | 0.909457 | 0.685351 | 00:00 |
| 2 | 0.684818 | 0.666282 | 00:00 |
| 3 | 0.593794 | 0.659886 | 00:00 |
learn.predict(ratings.iloc[0])
(FloatItem 4.166797, tensor(4.1668), tensor(4.1668))