from fastai.collab import *
from fastai.gen_doc.nbdoc import *
from fastai.callbacks.tensorboard import *

%load_ext tensorboard

show_doc(LearnerTensorboardWriter)

path = untar_data(URLs.ML_SAMPLE)

ratings = pd.read_csv(path/'ratings.csv')
series2cat(ratings, 'userId', 'movieId')
data = CollabDataBunch.from_df(ratings, seed=42)

learn = collab_learner(data, n_factors=30, y_range = [0, 5.5])

project_id = 'projct1'
tboard_path = Path('data/tensorboard/' + project_id)
learn.callback_fns.append(partial(LearnerTensorboardWriter, 
                                    base_dir=tboard_path, 
                                    name='run1'))

%tensorboard --logdir=$tboard_path --port=6006

learn.fit(10)

show_doc(LearnerTensorboardWriter.on_train_begin)

show_doc(LearnerTensorboardWriter.on_batch_end)

show_doc(LearnerTensorboardWriter.on_backward_end)

show_doc(LearnerTensorboardWriter.on_epoch_end)