diff --git a/recommender.py b/recommender.py index 3e0bdae7326de44f073cdd7d39b7f4259015bf10..05858236d9769f33fc86ea9decead7c130af20d7 100644 --- a/recommender.py +++ b/recommender.py @@ -7,6 +7,7 @@ from collections import defaultdict # Third-party imports import numpy as np import pandas as pd +import matplotlib.pyplot as plt from scipy.stats import pearsonr from sklearn.ensemble import AdaBoostRegressor, GradientBoostingRegressor, RandomForestRegressor @@ -23,7 +24,6 @@ from sklearn.svm import SVR from sklearn.tree import DecisionTreeRegressor from surprise import AlgoBase, KNNWithMeans, accuracy, PredictionImpossible from xgboost import XGBRegressor -from lightgbm import LGBMRegressor # Local imports from constants import Constant as C @@ -345,6 +345,19 @@ class RecommenderSystem_KNN : predictions = self.model.test(self.testset) self.rmse = accuracy.rmse(predictions) + def evaluate_rmse(self, testset): + """ + Evaluate the RMSE of the model on a test set. + + Args: + testset (list): List of testset entries containing (user, item, rating). + + Returns: + float: RMSE of the model. + """ + predictions = self.model.test(self.testset) + return accuracy.rmse(predictions, verbose=True) + def evaluate_mae(self): """ Evaluate the MAE of the model on the test set.