基於Python庫surprise的電影推薦系統
阿新 • • 發佈:2019-01-10
# -*- coding:utf-8 -*-
from __future__ import (absolute_import, division, print_function, unicode_literals)
import os
import io
from surprise import KNNBaseline
from surprise import Dataset
import logging
logging.basicConfig(level=logging.INFO,
format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s' ,
datefmt='%a, %d %b %Y %H:%M:%S')
# 訓練推薦模型 步驟:1
def getSimModle():
# 預設載入movielens資料集
data = Dataset.load_builtin('ml-100k')
trainset = data.build_full_trainset()
#使用pearson_baseline方式計算相似度 False以item為基準計算相似度 本例為電影之間的相似度
sim_options = {'name': 'pearson_baseline', 'user_based' : False}
##使用KNNBaseline演算法
algo = KNNBaseline(sim_options=sim_options)
#訓練模型
algo.train(trainset)
return algo
# 獲取id到name的互相對映 步驟:2
def read_item_names():
"""
獲取電影名到電影id 和 電影id到電影名的對映
"""
file_name = (os.path.expanduser('~') +
'/.surprise_data/ml-100k/ml-100k/u.item' )
rid_to_name = {}
name_to_rid = {}
with io.open(file_name, 'r', encoding='ISO-8859-1') as f:
for line in f:
line = line.split('|')
rid_to_name[line[0]] = line[1]
name_to_rid[line[1]] = line[0]
return rid_to_name, name_to_rid
# 基於之前訓練的模型 進行相關電影的推薦 步驟:3
def showSimilarMovies(algo, rid_to_name, name_to_rid):
# 獲得電影Toy Story (1995)的raw_id
toy_story_raw_id = name_to_rid['Toy Story (1995)']
logging.debug('raw_id=' + toy_story_raw_id)
#把電影的raw_id轉換為模型的內部id
toy_story_inner_id = algo.trainset.to_inner_iid(toy_story_raw_id)
logging.debug('inner_id=' + str(toy_story_inner_id))
#通過模型獲取推薦電影 這裡設定的是10部
toy_story_neighbors = algo.get_neighbors(toy_story_inner_id, 10)
logging.debug('neighbors_ids=' + str(toy_story_neighbors))
#模型內部id轉換為實際電影id
neighbors_raw_ids = [algo.trainset.to_raw_iid(inner_id) for inner_id in toy_story_neighbors]
#通過電影id列表 或得電影推薦列表
neighbors_movies = [rid_to_name[raw_id] for raw_id in neighbors_raw_ids]
print('The 10 nearest neighbors of Toy Story are:')
for movie in neighbors_movies:
print(movie)
if __name__ == '__main__':
# 獲取id到name的互相對映
rid_to_name, name_to_rid = read_item_names()
# 訓練推薦模型
algo = getSimModle()
##顯示相關電影
showSimilarMovies(algo, rid_to_name, name_to_rid)