-
Notifications
You must be signed in to change notification settings - Fork 1
/
Recommender.py
95 lines (79 loc) · 2.4 KB
/
Recommender.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import os
import numpy as np
import pandas as pd
import surprise
from surprise import Dataset, Reader, SVD
MODEL_NAME = "Surprise_SVD"
class Recommender:
"""
A class for Recommender
@author: Minyang Wang
@date: 9/14/2022
"""
def __init__(self):
"""
Initialize a SVD model
"""
self.model = SVD()
def fit(self, train_X, train_y):
"""
Fitting the model
Args:
train_X: X
train_y: y
"""
reader = Reader(rating_scale=(1, 5))
train = pd.DataFrame(train_X.copy())
train["rating"] = train_y
data = Dataset.load_from_df(train, reader)
self.model.fit(data.build_full_trainset())
def predict(self, X):
"""
Predict, given user-item pair
Args:
X: pd.DataFrame with shape (N, 2); first col = user_id, second col = item_id
"""
test = pd.DataFrame(X)
test.columns = ['u', 'i']
test["rating"] = 1
test_ = Dataset.load_from_df(test, reader=Reader(rating_scale=(1, 5))).build_full_trainset()
test_set = test_.build_testset()
predictions = pd.DataFrame(self.model.test(test_set))
temp2 = pd.merge(test, predictions, left_on=['u', 'i'], right_on=['uid', 'iid'], how="left")
return np.array(temp2.est).reshape(-1, 1)
def save(self, model_path):
"""
Save to local
Args:
model_path: path to write to
"""
surprise.dump.dump(os.path.join(model_path, "model.save"), predictions=None, algo=self.model, verbose=1)
@classmethod
def load(cls, model_path):
"""
Helper for loading model
Args:
model_path: path to model
Returns:
model, with weights loaded
"""
mf = cls()
mf.model = surprise.dump.load(os.path.join(model_path, "model.save"))[1]
return mf
def save_model(model, model_path):
model.save(model_path)
def load_model(model_path):
"""
Helper for loading model
Args:
model_path: path to model
Returns:
model
"""
try:
model = Recommender.load(model_path)
except:
raise Exception(f'''Error loading the trained {MODEL_NAME} model.
Did you accidentally delete the model? Please redownload
the zip file. Or you can run RecommderTraining.ipynb.''')
return model