Fix bugs; Add more demonstration for execution steps; Add input tables

app.py

@@ -49,11 +49,12 @@ def get_key():
     ip = requests.get('https://checkip.amazonaws.com').text.strip()
     print(ip)
 
-    URL = "http://54.242.37.195:8080/api/predict"
-    # The springboard machine we built to protect the key, 20217 is the birthday of Tianbao's girlfriend
-    # we will only let the demo machine have the access to the keys
-
-    one_key = requests.post(url=URL, json={"data": "Hi, binder server. Give me a key!"}).json()['data'][0]
+    # URL = "http://54.242.37.195:8080/api/predict"
+    # # The springboard machine we built to protect the key, 20217 is the birthday of Tianbao's girlfriend
+    # # we will only let the demo machine have the access to the keys
+    #
+    # one_key = requests.post(url=URL, json={"data": "Hi, binder server. Give me a key!"}).json()['data'][0]
+    one_key = "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"
     return one_key
 
 
@@ -96,6 +97,20 @@ def generate_binder_program(_args, _generator, _data_item):
     return response_dict["0"][0][0]
 
 
+def remove_row_id(table):
+    new_table = {"header": [], "rows": []}
+    header: list = table['header']
+    rows = table['rows']
+
+    if not 'row_id' in header:
+        return table
+
+    new_table['header'] = header[1:]
+    new_table['rows'] = [row[1:] for row in rows]
+
+    return new_table
+
+
 # Set up
 import nltk
 
@@ -217,34 +232,48 @@ try:
     if selected_language == 'SQL':
         with open("tmp_for_vis/{}_tmp_for_vis_steps.txt".format(stamp), "r") as f:
             steps = json.load(f)
-        col1, col2, col3 = st.columns([4.7, 0.6, 4.7])
-        # col1.subheader('Steps')
-        # col3.subheader('Intermediate results')
         for i, step in enumerate(steps):
-            col1, _, _ = st.columns([4.7, 0.6, 4.7])
+            col1, _, _ = st.columns([7, 1, 2])
             with col1:
                 st.markdown(f'**Step#{i + 1}**')
-            col1, col2, col3 = st.columns([4.7, 0.6, 4.7])
+            if i == len(steps) - 1:
+                col1, col1_25, col1_5, col2, col3 = st.columns([5, 0.3, 1.7, 1, 2])
+            else:
+                col1, col1_25, col1_5, col2, col3 = st.columns([3, 0.3, 2.7, 1, 2])
             with col1:
                 st.markdown('```sql\n' + step + '\n```')
-                # st.markdown('```' + step + '```')
-                # with st.container():
-                #     st.write(step)
+            with col1_25:
+                st.markdown("on")
+            with col1_5:
+                if i == len(steps) - 1:
+                    st.metric(label="whole", value="Table", delta = None)
+                else:
+                    with open("tmp_for_vis/{}_result_step_{}_input.txt".format(stamp, i), "r") as f:
+                        sub_tables_input = json.load(f)
+                    for sub_table in sub_tables_input:
+                        sub_table_to_print = remove_row_id(sub_table)
+                        st.dataframe(pd.DataFrame(sub_table_to_print['rows'], columns=sub_table_to_print['header']))
             with col2:
                 st.markdown('$\\rightarrow$')
+                if i == len(steps) - 1:
+                    # The final step
+                    st.metric(label="+", value="{}".format(selected_language), delta="Interpreter")
+                else:
+                    st.metric(label="+", value="Codex", delta="10+ examples")
             with st.spinner('...'):
                 time.sleep(1)
             with open("tmp_for_vis/{}_result_step_{}.txt".format(stamp, i), "r") as f:
                 result_in_this_step = json.load(f)
             with col3:
                 if isinstance(result_in_this_step, Dict):
-                    rows = result_in_this_step["rows"]
-                    header = result_in_this_step["header"]
+
+                    rows = remove_row_id(result_in_this_step)["rows"]
+                    header = remove_row_id(result_in_this_step)["header"]
                     if isinstance(header, list):
                         for idx in range(len(header)):
                             if header[idx].startswith('col_'):
                                 header[idx] = step
-                    st.dataframe(pd.DataFrame(pd.DataFrame(rows, columns=header)), use_container_width=True)
+                    st.dataframe(pd.DataFrame(rows, columns=header), use_container_width=True)
                 else:
                     st.markdown(result_in_this_step)
             with st.spinner('...'):

nsql/nsql_exec.py

@@ -120,6 +120,8 @@ class NSQLExecutor(object):
                                                            qa_type="map",
                                                            new_col_name_s=step.produced_col_name_s,
                                                            verbose=verbose)
+                        with open("tmp_for_vis/{}_result_step_{}_input.txt".format(stamp, steps.index(step)), "w") as f:
+                            json.dump(sql_executed_sub_tables, f)
                         with open("tmp_for_vis/{}_result_step_{}.txt".format(stamp, steps.index(step)), "w") as f:
                             json.dump(sub_table, f)
                         db.add_sub_table(sub_table, verbose=verbose)
@@ -131,6 +133,8 @@ class NSQLExecutor(object):
                                                            qa_type="map",
                                                            new_col_name_s=["col_{}".format(col_idx)],
                                                            verbose=verbose)
+                        with open("tmp_for_vis/{}_result_step_{}_input.txt".format(stamp, steps.index(step)), "w") as f:
+                            json.dump(sql_executed_sub_tables, f)
                         with open("tmp_for_vis/{}_result_step_{}.txt".format(stamp, steps.index(step)), "w") as f:
                             json.dump(sub_table, f)
                         return extract_answers(sub_table)
@@ -143,6 +147,8 @@ class NSQLExecutor(object):
                                                     table_title=db.table_title,
                                                     qa_type="ans",
                                                     verbose=verbose)
+                    with open("tmp_for_vis/{}_result_step_{}_input.txt".format(stamp, steps.index(step)), "w") as f:
+                        json.dump(sql_executed_sub_tables, f)
                     with open("tmp_for_vis/{}_result_step_{}.txt".format(stamp, steps.index(step)), "w") as f:
                         json.dump(answer, f)
                     if step.father: