Add total count to output

tapas_visualizer.py

@@ -126,10 +126,10 @@ class TapasVisualizer:
             cnt_html = (
                 f'<td style="border: none;" align="right">'
                 f'{self.style_span(str(cumulative_cnt), ["non-token", "count"])}'
-                '</td>'
+                "</td>"
                 f'<td style="border: none;" align="right">'
                 f'{self.style_span(f"<+{row_token_cnt}", ["non-token", "count"])}'
-                '</td>'
+                "</td>"
             )
             row_html = cnt_html + row_html
             table_html += f"<tr>{row_html}</tr>"
@@ -150,17 +150,13 @@ class TapasVisualizer:
             if int(segment_id) == 1:
                 cell_tokens[(row_id, col_id)].append(token_text)
 
-        table_html = (
-            '<tr><td style="border: none;" colspan="2" align="left">#Tokens</td></tr>'
-        )
-
         table_html, cumulative_cnt = self.cells_to_html(
             cell_vals=[table.columns],
             cell_tokens=cell_tokens,
             row_id_start=0,
             cell_element="th",
             cumulative_cnt=0,
-            table_html=table_html,
+            table_html="",
         )
 
         table_html, cumulative_cnt = self.cells_to_html(
@@ -171,6 +167,16 @@ class TapasVisualizer:
             cumulative_cnt=cumulative_cnt,
             table_html=table_html,
         )
+        top_label = self.style_span("#Tokens", ["count"])
+        top_label_cnt = self.style_span(f"(Total: {cumulative_cnt})", ["count"])
+
+        table_html = (
+            '<tr style="line-height: 2rem">'
+            f'<td style="border: none;" colspan="2" align="left">{top_label}</td>'
+            f'<td style="border: none;" colspan="1" align="left">{top_label_cnt}</td>'
+            "</tr>"
+            f"{table_html}"
+        )
 
         table_html = f"<table>{table_html}</table>"
         return HTMLBody(table_html)