ml100k_handler.cpp

#include "util/Base.h"
#include "struct/User.h"
#include "struct/Item.h"
#include "struct/Rating.h"
#include "util/FileUtil.h"

using namespace mf;

void transformIds(const string inputPath, unordered_map<string, int> &userIdMap, unordered_map<string, int> &itemIdMap) {

    string line;
    ifstream fin;
    fin.open(inputPath.c_str());

    vector<string> par;

    while (getline(fin, line)) {

        if (line.length() == 0) {
            continue;
        }
        par.clear();
        boost::split(par, line, boost::is_any_of("\t"));

        auto itr = userIdMap.find(par[0]);
        if (itr == userIdMap.end()) {
            userIdMap.insert(std::make_pair(par[0], userIdMap.size()));
        }

        itr = itemIdMap.find(par[1]);
        if (itr == itemIdMap.end()) {
            itemIdMap.insert(std::make_pair(par[1], itemIdMap.size()));
        }
    }
    fin.close();

}

void split(const string dataPath, const string outputFolder, const string metaPath, const string userIDMapPath, const string itemIDMapPath, const string outputMMCTrainPath,
           const string outputMMCTestPath, const string outputCSRTrainPath,
           const string outputCSRTestPath, const double percentage, unordered_map<string, int> &userIdMap, unordered_map<string, int> &itemIdMap) {

    string line;
    vector<User> users(userIdMap.size());
    vector<Item> items(itemIdMap.size());

    vector<string> par;

    int ratingNum = 0;

    ifstream fin(dataPath.c_str());
    while (getline(fin, line)) {
        if (line.length() == 0) {
            continue;
        }
        par.clear();
        boost::split(par, line, boost::is_any_of("\t"));

        int userID = userIdMap.find(par[0])->second;
        int itemID = itemIdMap.find(par[1])->second;
        users[userID].userID = userID;
        users[userID].ratedItemIDs.push_back(itemID);
        users[userID].ratings[itemID] = par[2];

        items[itemID].itemID = itemID;
        items[itemID].raterIDs.push_back(userID);
        items[itemID].ratings[userID] = par[2];

        ratingNum++;
    }
    fin.close();

    int trainNum = 0;
    int testNum = 0;

    vector<Rating> trainRatings;
    vector<Rating> testRatings;

    // each user has at least rated 20 items in the original data
    // first guarantee that each item at least appears in train set once
    for (int i = 0; i < items.size(); i++) {
        Item &item = items[i];
        unordered_map<int, string> &ratings = item.ratings;
        vector<int> &raterIDs = item.raterIDs;
        random_shuffle(raterIDs.begin(), raterIDs.end());

        int itemID = item.itemID;
        int userID = raterIDs[0];
        string score = ratings[userID];

        trainRatings.push_back(Rating(userID, itemID, ratings[userID]));
        trainNum++;

        //remove from users
        User &user = users[userID];
        user.ratedItemIDs.erase(std::remove(user.ratedItemIDs.begin(), user.ratedItemIDs.end(), itemID),
                                user.ratedItemIDs.end());
        user.ratings.erase(itemID);

    }

    double newPercentage = (ceil(ratingNum * percentage) - trainNum) / (ratingNum - trainNum + 0.0);

    for (User &user:users) {
        vector<int> &ratedItemIDs = user.ratedItemIDs;
        unordered_map<int, string> &ratings = user.ratings;
        random_shuffle(ratedItemIDs.begin(), ratedItemIDs.end());
        int flag = ceil(ratedItemIDs.size() * newPercentage);

        for (int i = 0; i < flag; i++) {
            trainRatings.push_back(Rating(user.userID, ratedItemIDs[i], ratings[ratedItemIDs[i]]));
            trainNum++;
        }

        for (int i = flag; i < user.ratings.size(); i++) {
            testRatings.push_back(Rating(user.userID, ratedItemIDs[i], ratings[ratedItemIDs[i]]));
            testNum++;
        }
    }

    writeMeta(outputFolder + metaPath, userIdMap.size(), itemIdMap.size(), trainRatings.size(), testRatings.size(), outputCSRTrainPath, outputCSRTestPath);

    writeMatrix(trainRatings, userIdMap.size(), itemIdMap.size(), outputFolder, outputMMCTrainPath, outputCSRTrainPath, "R_train_");

    writeMatrix(testRatings, userIdMap.size(), itemIdMap.size(), outputFolder, outputMMCTestPath, outputCSRTestPath, "R_test_");

    writeKeyMap(outputFolder + userIDMapPath, userIdMap);

    writeKeyMap(outputFolder + itemIDMapPath, itemIdMap);

}


int main(int argc, char const *argv[]){

    string outputFolder;
    string filePath;
    string metaPath;
    string userIDMapPath;
    string itemIDMapPath;
    string outputMMCTrainPath;
    string outputMMCTestPath;
    string outputCSRTrainPath;
    string outputCSRTestPath;

    // 80% of each user's ratings are used for training
    double percentage;

    po::options_description desc("Allowed options");
    desc.add_options()
            ("help", "produce help message")
            ("percentage", po::value<double>(&percentage)->default_value(0.8), "percentage of training ratings")
            ("rating_path", po::value<string>(&filePath)->default_value("../raw_data/ml-100k/u.data"), "path to original rating file")
            ("o_folder", po::value<string>(&outputFolder)->default_value("../data/ml-100k/"), "path to output folder")
            ("meta_path", po::value<string>(&metaPath)->default_value("meta"), "name of meta file")
            ("user_id_map_path", po::value<string>(&userIDMapPath)->default_value("user_id_map.dat"), "name of user id map file")
            ("item_id_map_path", po::value<string>(&itemIDMapPath)->default_value("item_id_map.dat"), "name of item id map file")
            ("o_mmc_train", po::value<string>(&outputMMCTrainPath)->default_value("train.mmc"), "name of output MMC training file")
            ("o_mmc_test", po::value<string>(&outputMMCTestPath)->default_value("test.mmc"), "name of output MMC testing file")
            ("o_csr_train", po::value<string>(&outputCSRTrainPath)->default_value("train.csr"), "name of output CSR training file")
            ("o_csr_test", po::value<string>(&outputCSRTestPath)->default_value("test.csr"), "name of output CSR testing file");

    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);

    if (vm.count("help")) {
        cout << desc << endl;
        return 0;
    }

    srand(time(NULL));

    unordered_map<string, int> userIdMap;
    unordered_map<string, int> itemIdMap;
    transformIds(filePath, userIdMap, itemIdMap);

    split(filePath, outputFolder, metaPath, userIDMapPath, itemIDMapPath, outputMMCTrainPath, outputMMCTestPath, outputCSRTrainPath,
          outputCSRTestPath, percentage, userIdMap, itemIdMap);

    return 0;
}