From 36c0816a61a070a0beedff63cc02e8c0536ba9d8 Mon Sep 17 00:00:00 2001 From: yifeim Date: Sun, 9 Feb 2020 16:48:55 -0800 Subject: [PATCH] Add an option to SAVE_TIMESTAMP `SAVE_TIMESTAMP=1 python netflix_data_convert.py {path/to/source} {path/to/target}` --- data_utils/netflix_data_convert.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/data_utils/netflix_data_convert.py b/data_utils/netflix_data_convert.py index 36d66f2..2293ecb 100644 --- a/data_utils/netflix_data_convert.py +++ b/data_utils/netflix_data_convert.py @@ -4,6 +4,7 @@ import sys import time import datetime +import os def print_stats(data): total_ratings = 0 @@ -17,7 +18,10 @@ def save_data_to_file(data, filename): with open(filename, 'w') as out: for userId in data: for record in data[userId]: - out.write("{}\t{}\t{}\n".format(userId, record[0], record[1])) + if not int(os.environ.get('SAVE_TIMESTAMP','0')): + out.write("{}\t{}\t{}\n".format(userId, record[0], record[1])) + else: + out.write("{}\t{}\t{}\t{}\n".format(userId, record[0], record[1], record[2])) def create_NETFLIX_data_timesplit(all_data, train_min,