def surprise_trainset_to_df(trainset,
col_user="******",
col_item="iid",
col_rating="rating"):
"""Converts a `surprise.Trainset` object to `pd.DataFrame`
More info: https://surprise.readthedocs.io/en/stable/trainset.html
Args:
trainset (obj): A surprise.Trainset object.
col_user (str): User column name.
col_item (str): Item column name.
col_rating (str): Rating column name.
Returns:
pd.DataFrame: A dataframe. The user and item columns are strings and the rating columns are floats.
"""
df = pd.DataFrame(trainset.all_ratings(),
columns=[col_user, col_item, col_rating])
map_user = (trainset._inner2raw_id_users
if trainset._inner2raw_id_users is not None else
invert_dictionary(trainset._raw2inner_id_users))
map_item = (trainset._inner2raw_id_items
if trainset._inner2raw_id_items is not None else
invert_dictionary(trainset._raw2inner_id_items))
df[col_user] = df[col_user].map(map_user)
df[col_item] = df[col_item].map(map_item)
return df
def surprise_trainset_to_df(
trainset, col_user="******", col_item="iid", col_rating="rating"
):
"""Converts a surprise.Trainset object to pd.DataFrame
More info: https://surprise.readthedocs.io/en/stable/trainset.html
Args:
trainset (obj): A surprise.Trainset object.
col_user (str): User column name.
col_item (str): Item column name.
col_rating (str): Rating column name.
Returns:
pd.DataFrame: A dataframe. The user and item columns are strings and the rating columns are floats.
"""
df = pd.DataFrame(trainset.all_ratings(), columns=[col_user, col_item, col_rating])
map_user = (
trainset._inner2raw_id_users
if trainset._inner2raw_id_users is not None
else invert_dictionary(trainset._raw2inner_id_users)
)
map_item = (
trainset._inner2raw_id_items
if trainset._inner2raw_id_items is not None
else invert_dictionary(trainset._raw2inner_id_items)
)
df[col_user] = df[col_user].map(map_user)
df[col_item] = df[col_item].map(map_item)
return df
def test_invert_dictionary():
d = {"a": 1, "b": 2}
d_inv = invert_dictionary(d)
assert d_inv == {1: "a", 2: "b"}
def test_invert_dictionary():
d = {"a": 1, "b": 2}
d_inv = invert_dictionary(d)
assert d_inv == {1: "a", 2: "b"}