36 lines
1.0 KiB
Python
36 lines
1.0 KiB
Python
import pytest
|
|
|
|
from pygrex.config import cfg
|
|
from pygrex.data_reader.data_reader import DataReader
|
|
from pygrex.models.als_model import ALS
|
|
from pygrex.models.bpr_model import BPR
|
|
from pygrex.recommender import Recommender
|
|
|
|
|
|
@pytest.fixture
|
|
def setup_data():
|
|
als = ALS(**cfg.model.als)
|
|
bpr = BPR(**cfg.model.bpr)
|
|
data = DataReader(**cfg.data.testdata)
|
|
data.make_consecutive_ids_in_dataset()
|
|
data.binarize()
|
|
return als, bpr, data
|
|
|
|
|
|
def test_train_als(setup_data):
|
|
als, _, data = setup_data
|
|
als.fit(data)
|
|
recommender = Recommender(data, als)
|
|
max_valid = als.model.item_factors.shape[0]
|
|
item_pool = list(range(min(max_valid, 50)))
|
|
_ = recommender.recommend(user_id=0, target_item_id=item_pool)
|
|
|
|
|
|
def test_train_bpr(setup_data):
|
|
_, bpr, data = setup_data
|
|
bpr.fit(data)
|
|
recommender = Recommender(data, bpr)
|
|
max_valid = bpr.model.item_factors.shape[0]
|
|
item_pool = list(range(min(max_valid, 50)))
|
|
_ = recommender.recommend(user_id=0, target_item_id=item_pool)
|