data_manager.py

'''
Created on Nov 9, 2015

@author: donghyun
'''

import os
import sys
import csv
import cPickle as pickl
import numpy as np
import glob
import ntpath

from operator import itemgetter
from scipy.sparse.csr import csr_matrix

from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer

import random


class Data_Factory():

    def load(self, path):
        R = pickl.load(open(path + "/ratings.all", "rb"))
        print "Load preprocessed rating data - %s" % (path + "/ratings.all")
        D_all = pickl.load(open(path + "/document.all", "rb"))
        print "Load preprocessed document data - %s" % (path + "/document.all")
        return R, D_all

    def load_ratings(self, path):
        R = pickl.load(open(path + "/ratings.all", "rb"))
        print "Load preprocessed rating data - %s" % (path + "/ratings.all")
        return R

    def load_documents(self, path):
        D_all = pickl.load(open(path + "/document.all", "rb"))
        print "Load preprocessed document data - %s" % (path + "/document.all")
        return D_all

    def save(self, path, R, D_all):
        if not os.path.exists(path):
            os.makedirs(path)
        print "Saving preprocessed rating data - %s" % (path + "/ratings.all")
        pickl.dump(R, open(path + "/ratings.all", "wb"))
        print "Done!"
        print "Saving preprocessed document data - %s" % (path + "/document.all")
        pickl.dump(D_all, open(path + "/document.all", "wb"))
        print "Done!"

    def read_rating(self, path):
        results = []
        if os.path.isfile(path):
            raw_ratings = open(path, 'r')
        else:
            print "Path (preprocessed) is wrong!"
            sys.exit()
        index_list = []
        rating_list = []
        all_line = raw_ratings.read().splitlines()
        for line in all_line:
            tmp = line.split()
            num_rating = int(tmp[0])
            if num_rating > 0:
                tmp_i, tmp_r = zip(*(elem.split(":") for elem in tmp[1::]))
                index_list.append(np.array(tmp_i, dtype=int))
                rating_list.append(np.array(tmp_r, dtype=float))
            else:
                index_list.append(np.array([], dtype=int))
                rating_list.append(np.array([], dtype=float))

        results.append(index_list)
        results.append(rating_list)

        return results

    def read_pretrained_word2vec(self, path, vocab, dim):
        if os.path.isfile(path):
            raw_word2vec = open(path, 'r')
        else:
            print "Path (word2vec) is wrong!"
            sys.exit()

        word2vec_dic = {}
        all_line = raw_word2vec.read().splitlines()
        mean = np.zeros(dim)
        count = 0
        for line in all_line:
            tmp = line.split()
            _word = tmp[0]
            _vec = np.array(tmp[1:], dtype=float)
            if _vec.shape[0] != dim:
                print "Mismatch the dimension of pre-trained word vector with word embedding dimension!"
                sys.exit()
            word2vec_dic[_word] = _vec
            mean = mean + _vec
            count = count + 1

        mean = mean / count

        W = np.zeros((len(vocab) + 1, dim))
        count = 0
        for _word, i in vocab:
            if word2vec_dic.has_key(_word):
                W[i + 1] = word2vec_dic[_word]
                count = count + 1
            else:
                W[i + 1] = np.random.normal(mean, 0.1, size=dim)

        print "%d words exist in the given pretrained model" % count

        return W

    def read_attributes(self, path):
        """
            Parse papers' features
        """
        if not os.path.isfile(path):
            print "Path (paper_info) is wrong!"
            sys.exit()
        else:
            print "Load preprocessed attributes data - %s" % (path)
            with open(path, "r") as f:
                reader = csv.reader(f, delimiter='\t')
                first_line = True
                feature_vec = []
                i = 0
                row_length = 0
                labels_ids = []
                for line in reader:
                    if first_line:
                        # labels = ['pages', 'year', 'type_NaN', 'type_article',
                        #           'type_book', 'type_booklet', 'type_electronic', 'type_inbook',
                        #           'type_incollection', 'type_inproceedings', 'type_manual',
                        #           'type_mastersthesis', 'type_misc', 'type_phdthesis', 'type_proceedings',
                        #           'type_techreport', 'type_unpublished']
                        labels = ['pages_norm', 'type_article', 'type_book', 'type_booklet',
                                   'type_electronic', 'type_inbook', 'type_incollection', 'type_inproceedings',
                                   'type_manual', 'type_mastersthesis', 'type_misc', 'type_phdthesis',
                                   'type_proceedings', 'type_techreport', 'type_unpublished', 'type_nan',
                                   'year_cat_null', 'year_cat_befor_90', 'year_cat_90-94', 'year_cat_95-1999',
                                   'year_cat_2000-2004', 'year_cat_2005-2016', 'year_cat_nan']
                        for j, entry in enumerate(line):
                            if entry in labels:
                                labels_ids.append(j)
                        row_length = len(labels_ids)
                        first_line = False
                        i += 1
                        continue
                    paper_id = line[0]
                    current_entry = []
                    for label_id in labels_ids:
                        current_entry.append(float(line[label_id]))
                    feature_vec.append(current_entry)
                    i += 1
            return labels, np.asarray(feature_vec)

    def split_data(self, ratio, R):
        print "Randomly splitting rating data into training set (%.1f) and test set (%.1f)..." % (1 - ratio, ratio)
        train = []  # make sure all users appear in the train set
        for i in xrange(R.shape[
                            0]):  # for each user, add the first item that he has rated, the resulted train list size is num_users
            user_rating = R[i].nonzero()[1]
            np.random.shuffle(user_rating)
            train.append((i, user_rating[0]))

        remain_item = set(xrange(R.shape[1])) - set(zip(*train)[1])  # get the item idx that haven't appeared in train.
        # make sure all items appear in the train set
        for j in remain_item:  # for each item in remain_item, get the list of users who rated this item, and add the if of the first user
            item_rating = R.tocsc().T[j].nonzero()[1]
            np.random.shuffle(item_rating)
            train.append((item_rating[0], j))  # the resulted list is of size num_user + len(remain_item)

        rating_list = set(zip(R.nonzero()[0], R.nonzero()[1]))
        total_size = len(rating_list)
        remain_rating_list = list(rating_list - set(train))
        random.shuffle(remain_rating_list)

        num_addition = int((1 - ratio) * total_size) - len(train)
        if num_addition < 0:
            print 'this ratio cannot be handled'
            sys.exit()
        else:
            train.extend(remain_rating_list[:num_addition])
            tmp_test = remain_rating_list[num_addition:]
            random.shuffle(tmp_test)
            valid = tmp_test[::2]
            test = tmp_test[1::2]

            trainset_u_idx, trainset_i_idx = zip(*train)
            trainset_u_idx = set(trainset_u_idx)
            trainset_i_idx = set(trainset_i_idx)
            if len(trainset_u_idx) != R.shape[0] or len(trainset_i_idx) != R.shape[1]:
                print "Fatal error in split function. Check your data again or contact authors"
                sys.exit()

        print "Finish constructing training set and test set"
        return train, valid, test

    # def read_split_from_file(self, ):

    def generate_train_valid_test_file_from_R(self, path, R, ratio):
        '''
        Split randomly rating matrix into training set, valid set and test set with given ratio (valid+test)
        and save three data sets to given path.
        Note that the training set contains at least a rating on every user and item.

        Input:
        - path: path to save training set, valid set, test set
        - R: rating matrix (csr_matrix)
        - ratio: (1-ratio), ratio/2 and ratio/2 of the entire dataset (R) will be training, valid and test set, respectively
        '''
        train, valid, test = self.split_data(ratio, R)
        print "Save training set and test set to %s..." % path
        if not os.path.exists(path):
            os.makedirs(path)

        R_lil = R.tolil()
        user_ratings_train = {}
        item_ratings_train = {}
        for i, j in train:
            if user_ratings_train.has_key(i):
                user_ratings_train[i].append(j)
            else:
                user_ratings_train[i] = [j]

            if item_ratings_train.has_key(j):
                item_ratings_train[j].append(i)
            else:
                item_ratings_train[j] = [i]

        user_ratings_valid = {}
        item_ratings_valid = {}
        for i, j in valid:
            if user_ratings_valid.has_key(i):
                user_ratings_valid[i].append(j)
            else:
                user_ratings_valid[i] = [j]

            if item_ratings_valid.has_key(j):
                item_ratings_valid[j].append(i)
            else:
                item_ratings_valid[j] = [i]

        user_ratings_test = {}
        item_ratings_test = {}
        for i, j in test:
            if user_ratings_test.has_key(i):
                user_ratings_test[i].append(j)
            else:
                user_ratings_test[i] = [j]

            if item_ratings_test.has_key(j):
                item_ratings_test[j].append(i)
            else:
                item_ratings_test[j] = [i]

        f_train_user = open(path + "/train_user.dat", "w")
        f_valid_user = open(path + "/valid_user.dat", "w")
        f_test_user = open(path + "/test_user.dat", "w")

        formatted_user_train = []
        formatted_user_valid = []
        formatted_user_test = []

        for i in xrange(R.shape[0]):
            if user_ratings_train.has_key(i):
                formatted = [str(len(user_ratings_train[i]))]
                formatted.extend(["%d:%.1f" % (j, R_lil[i, j])
                                  for j in sorted(user_ratings_train[i])])
                formatted_user_train.append(" ".join(formatted))
            else:
                formatted_user_train.append("0")

            if user_ratings_valid.has_key(i):
                formatted = [str(len(user_ratings_valid[i]))]
                formatted.extend(["%d:%.1f" % (j, R_lil[i, j])
                                  for j in sorted(user_ratings_valid[i])])
                formatted_user_valid.append(" ".join(formatted))
            else:
                formatted_user_valid.append("0")

            if user_ratings_test.has_key(i):
                formatted = [str(len(user_ratings_test[i]))]
                formatted.extend(["%d:%.1f" % (j, R_lil[i, j])
                                  for j in sorted(user_ratings_test[i])])
                formatted_user_test.append(" ".join(formatted))
            else:
                formatted_user_test.append("0")

        f_train_user.write("\n".join(formatted_user_train))
        f_valid_user.write("\n".join(formatted_user_valid))
        f_test_user.write("\n".join(formatted_user_test))

        f_train_user.close()
        f_valid_user.close()
        f_test_user.close()
        print "\ttrain_user.dat, valid_user.dat, test_user.dat files are generated."

        f_train_item = open(path + "/train_item.dat", "w")
        f_valid_item = open(path + "/valid_item.dat", "w")
        f_test_item = open(path + "/test_item.dat", "w")

        formatted_item_train = []
        formatted_item_valid = []
        formatted_item_test = []

        for j in xrange(R.shape[1]):
            if item_ratings_train.has_key(j):
                formatted = [str(len(item_ratings_train[j]))]
                formatted.extend(["%d:%.1f" % (i, R_lil[i, j])
                                  for i in sorted(item_ratings_train[j])])
                formatted_item_train.append(" ".join(formatted))
            else:
                formatted_item_train.append("0")

            if item_ratings_valid.has_key(j):
                formatted = [str(len(item_ratings_valid[j]))]
                formatted.extend(["%d:%.1f" % (i, R_lil[i, j])
                                  for i in sorted(item_ratings_valid[j])])
                formatted_item_valid.append(" ".join(formatted))
            else:
                formatted_item_valid.append("0")

            if item_ratings_test.has_key(j):
                formatted = [str(len(item_ratings_test[j]))]
                formatted.extend(["%d:%.1f" % (i, R_lil[i, j])
                                  for i in sorted(item_ratings_test[j])])
                formatted_item_test.append(" ".join(formatted))
            else:
                formatted_item_test.append("0")

        f_train_item.write("\n".join(formatted_item_train))
        f_valid_item.write("\n".join(formatted_item_valid))
        f_test_item.write("\n".join(formatted_item_test))

        f_train_item.close()
        f_valid_item.close()
        f_test_item.close()
        print "\ttrain_item.dat, valid_item.dat, test_item.dat files are generated."

        print "Done!"

    def generate_CTRCDLformat_content_file_from_D_all(self, path, D_all):
        '''
        Write word index with word count in document for CTR&CDL experiment

        '''
        f_text = open(path + "mult.dat", "w")
        X = D_all['X_base']
        formatted_text = []
        for i in xrange(X.shape[0]):
            word_count = sorted(set(X[i].nonzero()[1]))
            formatted = [str(len(word_count))]
            formatted.extend(["%d:%d" % (j, X[i, j]) for j in word_count])
            formatted_text.append(" ".join(formatted))

        f_text.write("\n".join(formatted_text))
        f_text.close()

    def generate_train_valid_test_from_ctr_split(self, inpath, outpath):

        if not os.path.exists(inpath):
            print "Path (Splits) is wrong! \n %s " % inpath
            sys.exit()
        print "Read training, test, and valid sets from %s..." % inpath

        print "Save training, test, and valid sets to %s..." % outpath
        if not os.path.exists(outpath):
            os.makedirs(outpath)

        f_train_user_in = glob.glob(os.path.join(inpath, "train-fold_*-users.dat"))[0]
        f_train_user_out = os.path.join(outpath, ntpath.basename(f_train_user_in))
        self.convert_and_save(f_train_user_in, f_train_user_out)

        if os.path.exists(os.path.join(inpath, "validation-fold_*-users.dat")):
            f_valid_user_in = glob.glob(os.path.join(inpath, "validation-fold_*-users.dat"))[0]
            f_valid_user_out = os.path.join(outpath, ntpath.basename(f_valid_user_in))
            self.convert_and_save(f_valid_user_in, f_valid_user_out)

        f_test_user_in = glob.glob(os.path.join(inpath, "test-fold_*-users.dat"))[0]
        f_test_user_out = os.path.join(outpath, ntpath.basename(f_test_user_in))
        self.convert_and_save(f_test_user_in, f_test_user_out)

        f_train_item_in = glob.glob(os.path.join(inpath, "train-fold_*-items.dat"))[0]
        f_train_item_out = os.path.join(outpath, ntpath.basename(f_train_item_in))
        self.convert_and_save(f_train_item_in, f_train_item_out)

        if os.path.exists(os.path.join(inpath, "validation-fold_*-items.dat")):
            f_valid_item_in = glob.glob(os.path.join(inpath, "validation-fold_*-items.dat"))[0]
            f_valid_item_out = os.path.join(outpath, ntpath.basename(f_valid_item_in))
            self.convert_and_save(f_valid_item_in, f_valid_item_out)

        f_test_item_in = glob.glob(os.path.join(inpath, "test-fold_*-items.dat"))[0]
        f_test_item_out = os.path.join(outpath, ntpath.basename(f_test_item_in))
        self.convert_and_save(f_test_item_in, f_test_item_out)

    def convert_and_save(self, infile, outfile):
        print('Reading ratings %s ...' % infile)
        formatted_list = []
        with open(infile, 'r') as f:
            with open(outfile, 'w') as of:
                row_id = 0  # 0 base index
                for line in f.readlines():
                    row_length = int(line.split()[0])
                    idx = map(int, line.split()[1:])
                    if row_length > 0:
                        if row_length != len(idx):
                            print "Row length doesn't match the number of elements "
                            sys.exit()
                        formatted = [str(row_length)]
                        formatted.extend(["%d:%.1f" % (i, 1.0)
                                          for i in sorted(idx)])
                        formatted_list.append(" ".join(formatted))
                    else:
                        formatted_list.append("0")
                of.write("\n".join(formatted_list))
        print('File {} is generated. \n'.format(outfile))

    def preprocess(self, path_rating, path_itemtext, min_rating,
                   _max_length, _max_df, _vocab_size):
        '''
        Preprocess rating and document data.

        Input:
            - path_rating: path for rating data (data format - user_id::item_id::rating)
            - path_itemtext: path for review or synopsis data (data format - item_id::text1|text2|text3|....)
            - min_rating: users who have less than "min_rating" ratings will be removed (default = 1)
            - _max_length: maximum length of document of each item (default = 300)
            - _max_df: terms will be ignored that have a document frequency higher than the given threshold (default = 0.5)
            - vocab_size: vocabulary size (default = 8000)

        Output:
            - R: rating matrix (csr_matrix: row - user, column - item)
            - D_all['X_sequence']: list of sequence of word index of each item ([[1,2,3,4,..],[2,3,4,...],...])
            - D_all['X_vocab']: list of tuple (word, index) in the given corpus
        '''
        # Validate data paths
        if os.path.isfile(path_rating):
            raw_ratings = open(path_rating, 'r')
            print "Path - rating data: %s" % path_rating
        else:
            print "Path(rating) is wrong!"
            sys.exit()

        if os.path.isfile(path_itemtext):
            raw_content = open(path_itemtext, 'r')
            print "Path - document data: %s" % path_itemtext
        else:
            print "Path(item text) is wrong!"
            sys.exit()

        # 1st scan document file to filter items which have documents
        tmp_id_plot = set()
        all_line = raw_content.read().splitlines()
        for line in all_line:
            tmp = line.split('::')
            i = tmp[0]
            tmp_plot = tmp[1].split('|')
            if tmp_plot[0] == '':
                continue
            tmp_id_plot.add(i)
        raw_content.close()

        print "Preprocessing rating data..."
        print "\tCounting # ratings of each user and removing users having less than %d ratings..." % min_rating
        # 1st scan rating file to check # ratings of each user
        all_line = raw_ratings.read().splitlines()
        tmp_user = {}
        for line in all_line:
            tmp = line.split('::')
            u = tmp[0]
            i = tmp[1]
            if (i in tmp_id_plot):
                if (u not in tmp_user):
                    tmp_user[u] = 1
                else:
                    tmp_user[u] = tmp_user[u] + 1

        raw_ratings.close()

        # 2nd scan rating file to make matrix indices of users and items
        # with removing users and items which are not satisfied with the given
        # condition
        raw_ratings = open(path_rating, 'r')
        all_line = raw_ratings.read().splitlines()
        userset = {}
        itemset = {}
        user_idx = 0
        item_idx = 0

        user = []
        item = []
        rating = []

        for line in all_line:
            tmp = line.split('::')
            u = tmp[0]
            # if u == '215':
            #     print u
            if u not in tmp_user:
                continue
            i = tmp[1]
            # An user will be skipped where the number of ratings of the user
            # is less than min_rating.
            if tmp_user[u] >= min_rating:
                if u not in userset:
                    userset[u] = user_idx
                    user_idx = user_idx + 1

                if (i not in itemset) and (i in tmp_id_plot):
                    itemset[i] = item_idx
                    item_idx = item_idx + 1
            else:
                continue

            if u in userset and i in itemset:
                u_idx = userset[u]
                i_idx = itemset[i]

                user.append(u_idx)
                item.append(i_idx)
                rating.append(float(tmp[2]))
            else:
                print '%s %s' % (u, i)

        raw_ratings.close()

        R = csr_matrix((rating, (user, item)))

        print "Finish preprocessing rating data - # user: %d, # item: %d, # ratings: %d" % (
        R.shape[0], R.shape[1], R.nnz)

        # 2nd scan document file to make idx2plot dictionary according to
        # indices of items in rating matrix
        print "Preprocessing item document..."

        # Read Document File
        raw_content = open(path_itemtext, 'r')
        max_length = _max_length
        map_idtoplot = {}
        all_line = raw_content.read().splitlines()
        for line in all_line:
            tmp = line.split('::')
            if tmp[0] in itemset:
                i = itemset[tmp[0]]
                tmp_plot = tmp[1].split('|')
                eachid_plot = (' '.join(tmp_plot)).split()[:max_length]
                map_idtoplot[i] = ' '.join(eachid_plot)

        print "\tRemoving stop words..."
        print "\tFiltering words by TF-IDF score with max_df: %.1f, vocab_size: %d" % (_max_df, _vocab_size)

        # Make vocabulary by document
        # vectorizer = TfidfVectorizer(max_df=_max_df, stop_words={
        #     'english'}, max_features=_vocab_size)
        # Make vocabulary by document
        vectorizer = TfidfVectorizer(max_df=_max_df, max_features=_vocab_size)
        Raw_X = [map_idtoplot[i] for i in range(R.shape[1])]
        vectorizer.fit(Raw_X)
        vocab = vectorizer.vocabulary_
        X_vocab = sorted(vocab.items(), key=itemgetter(1))

        # Make input for run
        X_sequence = []
        for i in range(R.shape[1]):
            X_sequence.append(
                [vocab[word] + 1 for word in map_idtoplot[i].split() if vocab.has_key(word)])

        '''Make input for CTR & CDL'''
        baseline_vectorizer = CountVectorizer(vocabulary=vocab)
        X_base = baseline_vectorizer.fit_transform(Raw_X)

        D_all = {
            'X_sequence': X_sequence,
            'X_base': X_base,
            'X_vocab': X_vocab,
        }

        print "Finish preprocessing document data!"

        return R, D_all