import sys import json from data_split import DataSplitter from create_vocab import HouseVocab import csv import time import argparse import pandas as pd from tqdm import tqdm import traceback count_saved = 0 master_count = 0 houses = [[] for _ in range(5)] # doing 5 houses at the moment # -------------------------------------------------------------------------- # Split Functions # -------------------------------------------------------------------------- def split_paragraph(paragraph): sentences = paragraph.split(".") # sentences = [sentence for sentence in sentences if not sentence.isdigit()] # sentences = [sentence + "." for sentence in sentences if len(sentence) > 0] sentences = [sentence + "." for sentence in sentences] return sentences def split_trajectory(trajectory): root = json.loads(trajectory) dios = root["dios"] npressed = root["nPressed"] sceneSwitches = root['sceneSwitches'] last_time = dios[-1]['time'] splits = [] current_split = [] counter = 0 for val in dios: if val["time"] > npressed[counter]: counter += 1 splits.append(current_split) current_split = [val] else: current_split.append(val) if len(current_split) > 0: splits.append(current_split) jsons_ = [] for i, split in enumerate(splits): prev_dios = [] for j in range(i): prev_dios.append(splits[j]) json_dict = {"dios": split, "prev_dios": prev_dios, "npressed": npressed, "sceneSwitches":sceneSwitches} jsons_.append(json_dict) return jsons_ # -------------------------------------------------------------------------- # Save Functions # -------------------------------------------------------------------------- def save_components(components, f): global count_saved global master_count try: sentences = components[4] traj = components[8] # for i, t in enumerate(traj): # dios = t['dios'] # last_point = dios[-1] # scene_switches = t['sceneSwitches'] # instruction = sentences[i] # debug = 1 if len(traj) > len(sentences) or len(traj) < len(sentences) - 1: return for sent_idx in range(min(len(sentences), len(traj))): for idx, comp in enumerate(components): # print(comp) if idx == 4: f.write(sentences[sent_idx] +'\n') # print('isntr: ' + sentences[sent_idx]) elif idx == 8: f.write(json.dumps(traj[sent_idx])) f.write('\n') else: f.write(str(comp) + '\n') if idx != 8: f.write('$$$$$$$$$$$$$$$$$$$$$$$$$$$\n') f.write('================================\n\n') count_saved += 1 except Exception: exc_info = sys.exc_info() traceback.print_exception(*exc_info) debug = 1 raise AssertionError("Failed.") # print('master count: ' + str(master_count)) # print('count saved: ' + str(count_saved)) def save_discrete_components(components): global master_count global count_saved # get house data house_name = components[6] sentences = components[4] traj = components[8] if 'house1' in house_name: house_index = 0 elif 'house2' in house_name: house_index = 1 elif 'house3' in house_name: house_index = 2 elif 'house4' in house_name: house_index = 3 elif 'house5' in house_name: house_index = 4 else: print('house index oops!') debug = 1 house_data = houses[house_index] # append new datapoints if len(traj) > len(sentences) or len(traj) < len(sentences) - 1: return for sent_idx in range(min(len(sentences), len(traj))): new_dp = {'ID': master_count, 'Instruction': sentences[sent_idx], 'House': house_name, 'Story': components[7], 'Trajectory': []} house_data.append(new_dp) master_count += 1 # print('master count: ' + str(master_count)) # print('count saved: ' + str(count_saved)) # -------------------------------------------------------------------------- # Parse Functinos # -------------------------------------------------------------------------- def parse_data_from_csv(from_csv_filename, to_txt_filename): def swap_components(a, b, components): components[a], components[b] = components[b], components[a] # swap order of components to match what was coming from txt file def reorder_components(components, df): components.remove(components[0]) components.remove(components[7]) components.remove(components[7]) swap_components(3, 1,components) swap_components(5, 1, components) swap_components(6, 2, components) swap_components(6, 4, components) swap_components(7, 6, components) story = df.loc[int(components[3])]['story'] components.insert(7, story) # read in story instruction data df = pd.DataFrame.from_csv('text2traj_instructions_2018_emnlp_instructions.csv') csv.field_size_limit(sys.maxsize) f_write = open(to_txt_filename, "w") with open(from_csv_filename, 'r') as csvfile: reader = csv.reader(csvfile) first_row = True for row in tqdm(reader): assert len(row) == 11 if first_row: first_row = False continue components = [] for i, el in enumerate(row): if i == 3: components.append(split_paragraph(el)) elif i == 5: components.append(split_trajectory(el)) else: components.append(el) reorder_components(components, df) save_components(components, f_write) save_discrete_components(components) f_write.flush() f_write.close() def parse_data_from_txt(from_txt_filename, to_txt_filename): lines = open(from_txt_filename).readlines() f_write = open(to_txt_filename, "w") # dataset = [] index = 0 components = [] sents, npressed = [], [] line_idx = -1 while line_idx < len(lines): line_idx += 1 if (line_idx % 1000 == 0): print(line_idx) ## sometimes a component is split amongst multiple lines curr_line, new_line = '', '' while (True): new_line = lines[line_idx].strip() new_datapoint = new_line.startswith("===========") new_component = new_line.startswith("$$$$$$$$$$$$") if new_component or new_datapoint: break curr_line += new_line line_idx += 1 if line_idx >= len(lines) -1: break if len(curr_line) == 0: if index < 1: continue if "Enter kitchen2" in curr_line: debug = 1 # check the current line and add to components if index == 4: # Instruction sentences = split_paragraph(curr_line) components.append(sentences) sents = sentences elif index == 8: # Trajectory traj = split_trajectory(curr_line) components.append(traj) root = json.loads(curr_line) dios = root["dios"] npressed = root["nPressed"] else: components.append(curr_line) ## check end of datapoint or end of component if new_line.startswith("==========="): # Add based on the components save_components(components, f_write) save_discrete_components(components) components = [] index = 0 continue if new_line.startswith("$$$$$$$$$$$$"): index += 1 continue f_write.flush() f_write.close() print("number segments saved: ", count_saved) print("number segments generates: ", master_count) return # -------------------------------------------------------------------------- # Create Functions Building # -------------------------------------------------------------------------- def create_discrete(houses): ## create house discrete files for house_idx, house_data in enumerate(houses): filename = "house" + str(house_idx+1) + "_discrete.json" json_obj = {'Dataset': house_data} with open(filename, 'w') as f: json.dump(json_obj, f) def create_vocab_and_split(houses): # vocab and split vocab_obj = HouseVocab(["house1","house2","house3", "house4","house5"]) vocab_filename = "house_all_vocab.txt" vocab_obj.create_vocab(vocab_filename) for house_idx, house_data in enumerate(houses): filename = "house" + str(house_idx + 1) + "_discrete.json" # Tokenize Instructions based on vocab vocab_obj.tokenize_instructions(house_idx) # split into train, dev, and test data splitter = DataSplitter('tokenized_'+filename) splitter.read_data() splitter.split_data([.75, .1, .15]) # -------------------------------------------------------------------------- # Main # -------------------------------------------------------------------------- def main(): # input arguments parser = argparse.ArgumentParser() parser.add_argument('--from_csv', type=bool, default=True) parser.add_argument('--from_txt', type=bool, default=False) parser.add_argument('--from_csv_filename', type=str, default='text2traj_trajectories_2018_emnlp.csv') parser.add_argument('--from_txt_filename', type=str, default='dataset_trajectory.txt') parser.add_argument('--to_txt_filename', type=str, default='dataset_trajectory_segments.txt') args = parser.parse_args() # parse input file and save dataset trajectory if args.from_txt: parse_data_from_txt(args.from_txt_filename, args.to_txt_filename) else: parse_data_from_csv(args.from_csv_filename, args.to_txt_filename) # Discretize the trajectories create_discrete(houses) # Build Vocab create_vocab_and_split(houses) if __name__ == "__main__": main()