def fit()

in src/rime/models/bpr.py [0:0]

• 34 lines of code
• 5 McCabe index (conditional complexity)

    def fit(self, D):
        ij_target = create_matrix(D.event_df, D.user_df.index, D.item_df.index, 'ij')
        dataset = np.array(ij_target, dtype=int).T

        N = len(dataset)
        train_set, valid_set = default_random_split(dataset)

        if 'embedding' in D.user_df:
            self._model_kw['user_embeddings'] = np.vstack(D.user_df['embedding'])
        if 'embedding' in D.item_df:
            self._model_kw['item_embeddings'] = np.vstack(D.item_df['embedding'])

        model = _BPR(len(D.user_df), len(D.item_df), **self._model_kw)
        model.user_proposal = (D.user_df['_hist_len'].values + 0.1) ** self.sample_with_posterior
        model.item_proposal = (D.item_df['_hist_len'].values + 0.1) ** self.sample_with_posterior

        if self.sample_with_prior:
            model.prior_score = auto_cast_lazy_score(getattr(D, "prior_score", None))
        else:
            model.prior_score = None

        trainer = Trainer(
            max_epochs=self.max_epochs, gpus=int(torch.cuda.is_available()),
            log_every_n_steps=1, callbacks=[model._checkpoint, LearningRateMonitor()])

        trainer.fit(
            model,
            DataLoader(train_set, self.batch_size, shuffle=True, num_workers=(N > 1e4) * 4),
            DataLoader(valid_set, self.batch_size, num_workers=(N > 1e4) * 4))
        model._load_best_checkpoint("best")
        for attr in ['user_proposal', 'item_proposal', 'prior_score', 'prior_score_T']:
            delattr(model, attr)

        self.D = D
        self.bpr_model = argparse.Namespace(
            user_embeddings=model.user_encoder.weight.detach().cpu().numpy(),
            item_embeddings=model.item_encoder.weight.detach().cpu().numpy(),
            user_biases=model.user_bias_vec.weight.squeeze(-1).detach().cpu().numpy(),
            item_biases=model.item_bias_vec.weight.squeeze(-1).detach().cpu().numpy(),
        )
        return self