def test_recommend_movies(pretrained, inputs, linear_feature_columns,
dnn_feature_columns, DEVICE):
# model = xDeepFM(linear_feature_columns, dnn_feature_columns, task='regression', device=DEVICE)
model = FiBiNET(linear_feature_columns,
dnn_feature_columns,
task='regression',
device=DEVICE)
model.load_state_dict(torch.load(pretrained))
pred_ans = model.predict(inputs, batch_size=256)
print(f'Predict rating: {pred_ans}')
pred_movie_list = []
idx = np.argsort(pred_ans, axis=0)[::-1]
for i, ans_idx in enumerate(idx[:, 0]):
if i < 5:
print(
f"Predict rating: {pred_ans[ans_idx][0] :.3f}, movie_id: {data.iloc[ans_idx]['movie_id']}, gender: {data.iloc[ans_idx]['gender']}, age: {data.iloc[ans_idx]['age']}, user_id: {data.iloc[ans_idx]['user_id']}"
)
pred_movie_list.append(data.iloc[ans_idx]['movie_id'])
# print(inputs['movie_id'].iloc[i])
# print('movie_max',data.loc[:, ['movie_id']].max(axis=0)) # 1682
return pred_movie_list
def recommend_movies(pretrained, inputs, linear_feature_columns,
dnn_feature_columns, DEVICE, df):
data = df
# print(data)
# model = xDeepFM(linear_feature_columns, dnn_feature_columns, task='regression', device=DEVICE)
model = FiBiNET(linear_feature_columns,
dnn_feature_columns,
task='regression',
device=DEVICE)
model.load_state_dict(torch.load(pretrained))
pred_ans = model.predict(inputs, batch_size=256)
# print(f'Predict rating: {pred_ans}')
pred_movie_list = []
pred_movie_genres = []
pred_rating = []
idx = np.argsort(pred_ans, axis=0)[::-1]
movie_genres = [
col for col in data.columns
if col not in ['movie_id', 'movie_title', 'unknown']
]
for i, ans_idx in enumerate(idx[:, 0]):
# TODO: Add arguments (rank)
if i < 2:
genres = [
movie_genre for movie_genre in movie_genres
if data.iloc[ans_idx][movie_genre] == 1
]
print(
f"Predict rating: {pred_ans[ans_idx][0] :.3f}, movie_id: {data.iloc[ans_idx]['movie_id']}, movie_title: {data.iloc[ans_idx]['movie_title']}, gender: {data.iloc[ans_idx]['gender']}, age: {data.iloc[ans_idx]['age']}, genres: {genres}"
)
pred_movie_list.append(data.iloc[ans_idx]['movie_title'])
pred_movie_genres.append(genres)
pred_rating.append(pred_ans[ans_idx][0])
# print(inputs['movie_id'].iloc[i])
# print('movie_max',data.loc[:, ['movie_id']].max(axis=0)) # 1682
return pred_movie_list, pred_movie_genres, pred_rating