版权声明:本文为博主原创文章,未经博主允许不得转载,欢迎大家讨论交流。 https://blog.csdn.net/qq_40006058/article/details/85389013
import os
from surprise import KNNBaseline
import io
from surprise import Dataset
# step 1 : train model
def TrainModel():
data = Dataset.load_builtin('ml-100k')
trainset = data.build_full_trainset()
# use pearson_baseline to compute similarity
sim_options = {'name' : 'pearson_baseline', 'user_based' : False}
algo = KNNBaseline(sim_options=sim_options)
# train
algo.fit(trainset)
return algo
# step 2 : get id_name and name_id
def Get_Dict():
file_name = os.path.expanduser('~/.surprise_data/ml-100k/ml-100k/u.item')
id_name = {}
name_id = {}
with open(file_name, 'r', encoding='ISO-8859-1') as f:
for line in f:
line = line.split('|')
id_name[line[0]] = line[1]
name_id[line[1]] = line[0]
return id_name, name_id
# step 3 : recommend movies based on the model
def RecommendMovie(movieName, algo, id_name, name_id, recommendNum):
# get movie's raw id
raw_id = name_id[movieName]
# translate raw_id to inner_id
inner_id = algo.trainset.to_inner_iid(raw_id)
# recommend movies
recommendations = algo.get_neighbors(inner_id, recommendNum)
# translate inner_id to raw_id
raw_ids = [algo.trainset.to_raw_iid(inner_id) for inner_id in recommendations]
# get movie name
movies = [id_name[raw_id] for raw_id in raw_ids]
for movie in movies:
print(movie)
if __name__ == '__main__':
id_name, name_id = Get_Dict()
algo = TrainModel()
showMovies = RecommendMovie('Toy Story (1995)', algo, id_name, name_id, 10)
Estimating biases using als…
Computing the pearson_baseline similarity matrix…
Done computing similarity matrix.
Beauty and the Beast (1991)
Raiders of the Lost Ark (1981)
That Thing You Do! (1996)
Lion King, The (1994)
Craft, The (1996)
Liar Liar (1997)
Aladdin (1992)
Cool Hand Luke (1967)
Winnie the Pooh and the Blustery Day (1968)
Indiana Jones and the Last Crusade (1989)