Refactor table_html creation

tapas_visualizer.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, List
+from typing import Any, List, Dict
 
 from collections import defaultdict
 
@@ -102,6 +102,31 @@ class TapasVisualizer:
 
         return spans
 
+    def cells_to_html(self,
+                      cell_vals: List[List[str]],
+                      cell_tokens: Dict,
+                      row_id_start: int=0,
+                      cell_element: str="td",
+                      cumulative_cnt: int=0,
+                      table_html: str="") -> str:
+
+        for row_id, row in enumerate(cell_vals, start=row_id_start):
+            row_html = ""
+            row_token_cnt = 0
+            for col_id, cell in enumerate(row, start=1):
+                cur_cell_tokens = cell_tokens[(row_id, col_id)]
+                span_htmls = self.text_to_html(cell, cur_cell_tokens)
+                cell_html = "".join(span_htmls)
+                row_html += f"<{cell_element}>{cell_html}</{cell_element}>"
+                row_token_cnt += len(cur_cell_tokens)
+            cumulative_cnt += row_token_cnt
+            # cnt_str = f'{row_token_cnt} | {cumulative_cnt}'
+            row_html += f'<td style="border: none;" align="right">{self.style_span(str(row_token_cnt), ["non-token", "count"])}</td>'
+            row_html += f'<td style="border: none;" align="right">{self.style_span(str(cumulative_cnt), ["non-token", "count"])}</td>'
+            table_html += f'<tr>{row_html}</tr>'
+
+        return table_html, cumulative_cnt
+
 
     def __call__(self, table: pd.DataFrame) -> Any:
         tokenized = self.tokenizer(table)
@@ -113,39 +138,22 @@ class TapasVisualizer:
             # 'prev_label', 'column_rank', 'inv_column_rank', 'numeric_relation' not required
             segment_id, col_id, row_id, *_ = tokenized['token_type_ids'][id_ind]
             token_text = self.tokenizer._convert_id_to_token(input_id)
-            cell_tokens[(row_id, col_id)].append(token_text)
-
-        cumulative_cnt = 0
-        header_row_html = ""
-        header_row_token_cnt = 0
-        for col_id, col in enumerate(table.columns, start=1):
-            cur_cell_tokens = cell_tokens[0, col_id]
-            span_htmls = self.text_to_html(col, cur_cell_tokens)
-            cell_html = "".join(span_htmls)
-            header_row_html += f"<th>{cell_html}</th>"
-            header_row_token_cnt += len(cur_cell_tokens)
-        cumulative_cnt += header_row_token_cnt
-        cnt_str = f'{header_row_token_cnt} | {cumulative_cnt}'
-        header_row_html += f'<th style="border: none;">{self.style_span(cnt_str, ["non-token", "count"])}</th>'
-        header_row_html = f'<tr>{header_row_html}</tr>'
-
-        table_vals = table.values
-
-        table_html = header_row_html
-
-        for row_id, row in enumerate(table_vals, start=1):
-            row_html = ""
-            row_token_cnt = 0
-            for col_id, cell in enumerate(row, start=1):
-                cur_cell_tokens = cell_tokens[(row_id, col_id)]
-                span_htmls = self.text_to_html(cell, cur_cell_tokens)
-                cell_html = "".join(span_htmls)
-                row_html += f"<td>{cell_html}</td>"
-                row_token_cnt += len(cur_cell_tokens)
-            cumulative_cnt += row_token_cnt
-            cnt_str = f'{row_token_cnt} | {cumulative_cnt}'
-            row_html += f'<td style="border: none;">{self.style_span(cnt_str, ["non-token", "count"])}</td>'
-            table_html += f'<tr>{row_html}</tr>'
+            if int(segment_id) == 1:
+                cell_tokens[(row_id, col_id)].append(token_text)
+
+        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, cumulative_cnt = self.cells_to_html(cell_vals=table.values,
+                                                        cell_tokens=cell_tokens,
+                                                        row_id_start=1,
+                                                        cell_element="td",
+                                                        cumulative_cnt=cumulative_cnt,
+                                                        table_html=table_html)
 
         table_html = f'<table>{table_html}</table>'
         return HTMLBody(table_html)