ift6758

#!/usr/bin/env python

import argparse
import os
import xml.etree.ElementTree as ET

from sklearn.externals import joblib

from utils.data_processing import parse_input, parse_output

model_path = "trained_models/final.pkl"
TEAM_NAME = "user17"


def evaluate(test_data_dir, results_output_dir):
    """
    Generates evaluation results as xml files
    :param test_data_dir: Root test data folder path
    :param results_output_dir: Output path to save xml results
    :return:
    """
    os.makedirs(results_output_dir, exist_ok=True)
    model = joblib.load(model_path)

    test_data = parse_input(test_data_dir, is_train=False)
    pred_df = model.predict(test_data)
    pred_df = parse_output(pred_df)

    for user_id in test_data['user_id']:
        pred = {}
        pred["id"] = str(user_id)
        pred.update(dict(pred_df.loc[user_id]))
        users_root = ET.Element('user', attrib=pred)
        xml_string_data = ET.tostring(users_root, encoding="unicode")
        xml_file = open(os.path.join(results_output_dir,
                                     "{}.xml".format(user_id)), "w")
        xml_file.write(xml_string_data)


if __name__ == "__main__":
    """
    Evaluation script that generates predictions and save them as xml files
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', type=str, default=None,
                        help='Data input directory')
    parser.add_argument('-o', type=str, default=None,
                        help='Output directory')
    args = parser.parse_args()
    evaluate(args.i, args.o)