def extract_actions()

in mm_action_prediction/tools/extract_actions.py [0:0]
72 lines of code
5 McCabe index (conditional complexity)

def extract_actions(input_json_file, save_root, furniture_db, subtask):
    """Extract assistant API calls from keystrokes and NLU/NLG annotations.

    Args:
        input_json_file: JSON data file
        save_root: Folder to save the extracted API calls
        furniture_db: object wrapping the furniture database
        subtask: Single dominant or multiple actions
    """
    # Read the raw data.
    print("Reading: {}".format(input_json_file))
    with open(input_json_file, "r") as file_id:
        data = json.load(file_id)

    dialogs = []
    price_dict = furniture_db.get_min_max_price_per_class()
    for datum in data["dialogue_data"]:
        dialog_id = datum["dialogue_idx"]
        dialog_datum = datum["dialogue"]
        dialog_coref_map = datum["dialogue_coref_map"]
        reversed_dialog_coref_map = {v: k for k, v in dialog_coref_map.items()}
        keystroke_sequence_with_args = []
        chat_utterances = []
        search_results = []
        last_search_args = {}

        for round_datum in dialog_datum:
            insert_item = {"turn_idx": round_datum["turn_idx"]}
            raw_keystroke_list = round_datum["raw_assistant_keystrokes"]
            keystrokes_with_args = get_keystrokes_with_args(
                raw_keystroke_list, price_dict
            )
            keystroke_sequence_with_args.extend(keystrokes_with_args)
            this_turn_keystrokes, keystroke_sequence_with_args = get_turn_keystrokes(
                keystroke_sequence_with_args
            )
            # get min. set of actions required to update the scene and
            # secondly furniture clicks which could signal viewed-text
            relevant_actions, viewed_text_keystrokes,\
                    search_results, last_search_args = \
                get_relevant_actions(
                    this_turn_keystrokes,
                    search_results,
                    last_search_args,
                    furniture_db
                )
            viewed_text_actions = get_viewed_text_actions(
                viewed_text_keystrokes
            )

            # get additional actions based on NLU/NLG annotation
            getinfo_actions = gen_getinfo_from_annotation(
                round_datum,
                reversed_dialog_coref_map
            )
            addtocart_actions = gen_addtocart_from_annotation(
                round_datum,
                reversed_dialog_coref_map
            )

            # collate the different sets of actions into insert_item
            collate_and_insert_actions(
                subtask,
                insert_item,
                relevant_actions,
                getinfo_actions,
                addtocart_actions,
                viewed_text_actions,
                round_datum
            )
            insert_item["raw_action_with_args"] = copy.deepcopy(keystrokes_with_args)
            insert_item["current_search_results"] = copy.deepcopy(search_results)
            chat_utterances.append(insert_item)

        roundwise_actions = get_roundwise_dialog_actions(
            subtask,
            chat_utterances
        )
        dialogs.append(
            {
                "dialog_id": dialog_id,
                "turns": chat_utterances,
                "actions": roundwise_actions,
            }
        )

    save_path = input_json_file.split("/")[-1].replace(".json", "_api_calls.json")
    save_path = os.path.join(save_root, save_path)
    print("Saving: {}".format(save_path))
    with open(save_path, "w") as f:
        json.dump(dialogs, f)