fix: support cp4d format (#102)

watson-developer-cloud · May 27, 2021 · 1e75b80 · 1e75b80
1 parent cb9eb02
commit 1e75b80
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/src/assistant_improve_toolkit/computation_func.py b/src/assistant_improve_toolkit/computation_func.py
@@ -300,16 +300,19 @@ def format_data(df):
     if 'response_context_response_context_IntentStarted' in df3.columns \
             and 'response_context_response_context_IntentCompleted' in df3.columns:
         cols = ['log_id', 'response_timestamp', 'response_context_conversation_id', 'request_input', 'response_text',
-                'response_intents', 'response_entities', 'response_nodes_visited', 'response_dialog_request_counter',
-                'response_dialog_turn_counter',
+                'response_intents', 'response_entities', 'response_nodes_visited', 'response_dialog_stack',
+                'response_dialog_request_counter', 'response_dialog_turn_counter',
                 'response_context_response_context_IntentStarted', 'response_context_response_context_IntentCompleted']
     else:
         cols = ['log_id', 'response_timestamp', 'response_context_conversation_id', 'request_input', 'response_text',
-                'response_intents', 'response_entities', 'response_nodes_visited', 'response_dialog_request_counter',
-                'response_dialog_turn_counter']
+                'response_intents', 'response_entities', 'response_nodes_visited', 'response_dialog_stack',
+                'response_dialog_request_counter', 'response_dialog_turn_counter']
 
-    if 'response_dialog_stack' in df3.columns:
-        cols.append('response_dialog_stack')
+    if 'response_nodes_visited' not in df3.columns:
+        df3['response_nodes_visited'] = [[] for _ in range(len(df3))]
+
+    if 'response_dialog_stack' not in df3.columns:
+        df3['response_dialog_stack'] = [[] for _ in range(len(df3))]
 
     print('Extracting intents ...')
     # Select a few required columns