Add/update docstring

README.md

@@ -1,6 +1,6 @@
 ---
 title: Tapas Tokenizer Viz
-emoji: 🌍
+emoji: 🍽️
 colorFrom: blue
 colorTo: indigo
 sdk: streamlit

tapas_visualizer.py

@@ -1,3 +1,8 @@
+"""Visualizer for TAPAS
+
+Implementation heavily based on
+`EncodingVisualizer` from `tokenizers.tools`.
+"""
 import os
 from typing import Any, List, Dict
 
@@ -5,22 +10,24 @@ from collections import defaultdict
 
 import pandas as pd
 
+from transformers import TapasTokenizer
+
 dirname = os.path.dirname(__file__)
 css_filename = os.path.join(dirname, "tapas-styles.css")
 with open(css_filename) as f:
     css = f.read()
 
 
-def HTMLBody(table_html: str, css_styles=css) -> str:
+def HTMLBody(table_html: str, css_styles: str = css) -> str:
     """
     Generates the full html with css from a list of html spans
 
     Args:
-        children (:obj:`List[str]`):
-            A list of strings, assumed to be html elements
+        table_html (str):
+            The html string of the table
 
-        css_styles (:obj:`str`, `optional`):
-            Optional alternative implementation of the css
+        css_styles (str):
+            CSS styling to be embedded inline
 
     Returns:
         :obj:`str`: An HTML string with style markup
@@ -42,10 +49,11 @@ def HTMLBody(table_html: str, css_styles=css) -> str:
 
 
 class TapasVisualizer:
-    def __init__(self, tokenizer) -> None:
+    def __init__(self, tokenizer: TapasTokenizer) -> None:
         self.tokenizer = tokenizer
 
     def normalize_token_str(self, token_str: str) -> str:
+        # Normalize subword tokens to org subword str
         return token_str.replace("##", "")
 
     def style_span(self, span_text: str, css_classes: List[str]) -> str: