RecBoleを使ってみよう5　推論

前回は、モデルの学習方法について解説しました。

今回は、次のテストデータについて、学習済みモデルのレコメンド先を推論します。

test.csv

user_id
5
1
2222
3333

最終コード

まずコードの全体を載せます。
テストファイルは ./data/processed/anime/test.csv、モデルファイルは ./saved/BPR-Sep-15-2023_16-55-36.pth に入っているとします。

predict.py

./predict.py

from __future__ import annotations

import logging

import numpy as np
import pandas as pd
import torch
from recbole.data.dataloader.general_dataloader import FullSortEvalDataLoader
from recbole.data.dataset.dataset import Dataset
from recbole.model.abstract_recommender import AbstractRecommender
from recbole.quick_start import load_data_and_model
from recbole.utils.case_study import full_sort_topk

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s [%(levelname)s]: %(message)s"
)


def full_sort_topk_wrapper(
    dataset: Dataset,
    external_user_ids: np.ndarray,
    model: AbstractRecommender,
    test_data: FullSortEvalDataLoader,
    topk: int,
) -> tuple[np.ndarray]:
    """
    指定されたデータセットとレコメンドモデルを用いて、各ユーザーに対するトップkのレコメンド先を生成します。
    フルソート評価を行います。

    Parameters:
    -----------
    dataset : Dataset
        設定とマッピング情報を含むデータセットオブジェクト。

    external_user_ids : np.ndarray
        レコメンドを生成する対象となる外部ユーザーIDの配列。

    model : AbstractRecommender
        レコメンドを生成するためのレコメンドモデル。

    test_data : FullSortEvalDataLoader
        テストデータを含む評価用のデータローダー。

    topk : int
        各ユーザーに対してレコメンドする上位k個のアイテムの数。

    Returns:
    --------
    tuple[np.ndarray]
        2つのNumPy配列を含むタプル:
        1) 各ユーザーに対する上位kのレコメンドアイテムのスコア。
        2) 各ユーザーに対する上位kのレコメンドアイテムの外部ID。
    """
    user_id_field = dataset.config.USER_ID_FIELD
    item_id_field = dataset.config.ITEM_ID_FIELD
    internal_user_ids: np.ndarray = dataset.token2id(
        field=user_id_field, tokens=external_user_ids
    )
    topk_scores_tensor: torch.Tensor
    topk_internal_item_ids_tensor: torch.Tensor
    topk_scores_tensor, topk_internal_item_ids_tensor = full_sort_topk(
        uid_series=internal_user_ids,
        model=model,
        test_data=test_data,
        k=topk,
        device=None,
    )
    topk_scores = topk_scores_tensor.cpu().numpy()
    topk_internal_item_ids = topk_internal_item_ids_tensor.cpu().numpy()
    topk_external_item_ids: np.ndarray = dataset.id2token(
        field=item_id_field, ids=topk_internal_item_ids
    )
    return topk_scores, topk_external_item_ids


def create_recommendation_df(
    external_user_ids: np.ndarray,
    topk_external_item_ids: np.ndarray,
    topk_scores: np.ndarray,
) -> pd.DataFrame:
    """
    各ユーザに対するトップkのレコメンドアイテムとスコアから、Pandas DataFrameを生成する。

    Parameters:
    -----------
    external_user_ids : np.ndarray (legnth: num_users)
        レコメンドを行ったユーザの外部ID

    topk_external_item_ids : np.ndarray (shape: (num_users, topk))
        各ユーザに対するトップkレコメンドアイテムの外部ID

    topk_scores : np.ndarray (shape: (num_users, topk))
        各ユーザに対するトップkレコメンドアイテムのスコア

    Returns:
    --------
    pd.DataFrame (shape: (num_users * topk, 4))
        user_id, rank（順位）, item_id, score をカラムに持つデータフレーム。
    """
    # 入力データのshapeチェック
    if (external_user_ids.shape[0] != topk_external_item_ids.shape[0]) or (
        external_user_ids.shape[0] != topk_scores.shape[0]
        or (topk_external_item_ids.shape[1] != topk_scores.shape[1])
    ):
        raise ValueError(
            f"入力データのshapeが想定されたものではありません\n"
            f"external_user_ids.shape: {external_user_ids.shape}\n"
            f"topk_external_item_ids.shape: {topk_external_item_ids.shape}\n"
            f"topk_scores.shape: {topk_scores.shape}"
        )

    # 入力データを加工してデータフレームを作成
    topk_scores_and_external_item_ids = np.dstack(
        (topk_scores, topk_external_item_ids)
    )  # shape: (num_users, topk, 2)
    recommendation_data = [
        {
            "user_id": external_user_ids[user_id_index],
            "rank": rank + 1,
            "score": float(score),
            "item_id": int(item_id),
        }
        for user_id_index, score_and_item_pair in enumerate(
            topk_scores_and_external_item_ids
        )
        for rank, (score, item_id) in enumerate(score_and_item_pair)
    ]
    return pd.DataFrame(recommendation_data)


def main():
    topk = 10
    test_file_path = "./data/processed/anime/test.csv"
    model_file_path = "./saved/BPR-Sep-15-2023_16-55-36.pth"
    output_file_path = "./data/output/recommendation_prediction.csv"

    # データとモデルと設定のロード
    config, model, dataset, train_data, valid_data, test_data = load_data_and_model(
        model_file=model_file_path,
    )

    # ユーザIDに使われてる物理名を取得（デフォルトは "user_id"）
    user_id_field = config["USER_ID_FIELD"]

    # テストデータを外部ユーザIDとして読み込む
    external_user_ids = np.array(
        pd.read_csv(
            test_file_path, usecols=[user_id_field], dtype={user_id_field: str}
        ).squeeze(axis=1)
    )

    # トップkのアイテムIDとスコアを算出
    topk_scores, topk_external_item_ids = full_sort_topk_wrapper(
        dataset=dataset,
        external_user_ids=external_user_ids,
        model=model,
        test_data=test_data,
        topk=topk,
    )

    # 上の結果をデータフレームにまとめる
    recommendation_df = create_recommendation_df(
        external_user_ids=external_user_ids,
        topk_external_item_ids=topk_external_item_ids,
        topk_scores=topk_scores,
    )

    # csv出力
    recommendation_df.to_csv(output_file_path, index=False)
    logging.info(f"File saved: {output_file_path}")


if __name__ == "__main__":
    main()

コマンド実行によって、推論結果が出力されます。

poetry run python predict.py

./data/output/recommendation_prediction.csv

user_id,rank,score,item_id
5,1,4.7843237,16498
5,2,4.7309036,5114
5,3,4.7251315,1575
5,4,4.6202517,226
5,5,4.611123,4224
5,6,4.6054974,2167
5,7,4.5779295,10620
5,8,4.551989,121
5,9,4.5149107,11111
5,10,4.5055494,2904
...（中略）...
3333,8,2.7181635,2167
3333,9,2.701967,2904
3333,10,2.684104,4224

コードの各部分の解説

TBA 余裕があったら加筆します！