import pandas as pd
from surprise import Dataset
from surprise import Reader
from surprise import SVD
from surprise import accuracy
from surprise.model_selection import train_test_split
data = Dataset.load_builtin('ml-100k')
df = pd.DataFrame(data.__dict__['raw_ratings'], columns = ['user_id', 'item_id', 'rating', 'timestamp'])
df.drop(['timestamp'], axis = 1, inplace = True)
reader = Reader(rating_scale = (1, 5))
data = Dataset.load_from_df(df, reader)
train_set, test_set = train_test_split(data, test_size = 0.25, random_state = 1004)
algo = SVD()
algo.fit(train_set)
pred = algo.test(test_set)
# accuracy.rmse(pred)
user_id = str(244)
item_id = str(51)
pred = algo.predict(user_id, item_id)
real = df[(df['user_id'] == user_id) & (df['item_id'] == item_id)]
print(f'예측 평점: {pred.est}')
print(f'실제 평점: {real.rating.iloc[0]}')
