codekingpro/portable-devtools
115k
1import html
2import itertools
3import os
4import re
5from string import Template
6from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple
7
8from tokenizers import Encoding, Tokenizer
9
10dirname = os.path.dirname(__file__)
11css_filename = os.path.join(dirname, "visualizer-styles.css")
12with open(css_filename) as f:
13 css = f.read()
14
15
16class Annotation:
17 start: int
18 end: int
19 label: str
20
21 def __init__(self, start: int, end: int, label: str):
22 self.start = start
23 self.end = end
24 self.label = label
25
26
27AnnotationList = List[Annotation]
28PartialIntList = List[Optional[int]]
29
30
31class CharStateKey(NamedTuple):
32 token_ix: Optional[int]
33 anno_ix: Optional[int]
34
35
36class CharState:
37 char_ix: Optional[int]
38
39 def __init__(self, char_ix):
40 self.char_ix = char_ix
41
42 self.anno_ix: Optional[int] = None
43 self.tokens: List[int] = []
44
45 @property
46 def token_ix(self):
47 return self.tokens[0] if len(self.tokens) > 0 else None
48
49 @property
50 def is_multitoken(self):
51 """
52 BPE tokenizers can output more than one token for a char
53 """
54 return len(self.tokens) > 1
55
56 def partition_key(self) -> CharStateKey:
57 return CharStateKey(
58 token_ix=self.token_ix,
59 anno_ix=self.anno_ix,
60 )
61
62
63class Aligned:
64 pass
65
66
67class EncodingVisualizer:
68 """
69 Build an EncodingVisualizer
70
71 Args:
72
73 tokenizer (:class:`~tokenizers.Tokenizer`):
74 A tokenizer instance
75
76 default_to_notebook (:obj:`bool`):
77 Whether to render html output in a notebook by default
78
79 annotation_converter (:obj:`Callable`, `optional`):
80 An optional (lambda) function that takes an annotation in any format and returns
81 an Annotation object
82 """
83
84 unk_token_regex = re.compile("(.{1}\b)?(unk|oov)(\b.{1})?", flags=re.IGNORECASE)
85
86 def __init__(
87 self,
88 tokenizer: Tokenizer,
89 default_to_notebook: bool = True,
90 annotation_converter: Optional[Callable[[Any], Annotation]] = None,
91 ):
92 if default_to_notebook:
93 try:
94 from IPython.display import HTML, display # type: ignore[attr-defined]
95 except ImportError:
96 try:
97 from IPython.core.display import HTML, display # type: ignore[attr-defined]
98 except ImportError:
99 msg = (
100 "We couldn't import IPython utils for html display.\n"
101 "Are you running in a notebook?\n"
102 "You can also pass `default_to_notebook=False` to get back raw HTML.\n"
103 )
104 raise ImportError(msg) from None
105 self.tokenizer = tokenizer
106 self.default_to_notebook = default_to_notebook
107 self.annotation_coverter = annotation_converter
108 pass
109
110 def __call__(
111 self,
112 text: str,
113 annotations: Optional[List[Any]] = None,
114 default_to_notebook: Optional[bool] = None,
115 ) -> Optional[str]:
116 """
117 Build a visualization of the given text
118
119 Args:
120 text (:obj:`str`):
121 The text to tokenize
122
123 annotations (:obj:`List[Annotation]`, `optional`):
124 An optional list of annotations of the text. The can either be an annotation class
125 or anything else if you instantiated the visualizer with a converter function
126
127 default_to_notebook (:obj:`bool`, `optional`, defaults to `False`):
128 If True, will render the html in a notebook. Otherwise returns an html string.
129
130 Returns:
131 The HTML string if default_to_notebook is False, otherwise (default) returns None and
132 renders the HTML in the notebook
133
134 """
135 final_default_to_notebook = self.default_to_notebook
136 if default_to_notebook is not None:
137 final_default_to_notebook = default_to_notebook
138 if final_default_to_notebook:
139 try:
140 from IPython.display import HTML, display # type: ignore[attr-defined]
141 except ImportError:
142 try:
143 from IPython.core.display import HTML, display # type: ignore[attr-defined]
144 except ImportError:
145 msg = (
146 "We couldn't import IPython utils for html display.\n"
147 "Are you running in a notebook?\n"
148 "You can also pass `default_to_notebook=False` to get back raw HTML.\n"
149 )
150 raise ImportError(msg) from None
151 if annotations is None:
152 annotations = []
153 if self.annotation_coverter is not None:
154 annotations = list(map(self.annotation_coverter, annotations))
155 encoding = self.tokenizer.encode(text)
156 html = EncodingVisualizer.__make_html(text, encoding, annotations)
157 if final_default_to_notebook:
158 display(HTML(html))
159 else:
160 return html
161
162 @staticmethod
163 def calculate_label_colors(annotations: AnnotationList) -> Dict[str, str]:
164 """
165 Generates a color palette for all the labels in a given set of annotations
166
167 Args:
168 annotations (:obj:`Annotation`):
169 A list of annotations
170
171 Returns:
172 :obj:`dict`: A dictionary mapping labels to colors in HSL format
173 """
174 if len(annotations) == 0:
175 return {}
176 labels = set(map(lambda x: x.label, annotations))
177 num_labels = len(labels)
178 h_step = int(255 / num_labels)
179 if h_step < 20:
180 h_step = 20
181 s = 32
182 l = 64 # noqa: E741
183 h = 10
184 colors = {}
185
186 for label in sorted(labels): # sort so we always get the same colors for a given set of labels
187 colors[label] = f"hsl({h},{s}%,{l}%)"
188 h += h_step
189 return colors
190
191 @staticmethod
192 def consecutive_chars_to_html(
193 consecutive_chars_list: List[CharState],
194 text: str,
195 encoding: Encoding,
196 ):
197 """
198 Converts a list of "consecutive chars" into a single HTML element.
199 Chars are consecutive if they fall under the same word, token and annotation.
200 The CharState class is a named tuple with a "partition_key" method that makes it easy to
201 compare if two chars are consecutive.
202
203 Args:
204 consecutive_chars_list (:obj:`List[CharState]`):
205 A list of CharStates that have been grouped together
206
207 text (:obj:`str`):
208 The original text being processed
209
210 encoding (:class:`~tokenizers.Encoding`):
211 The encoding returned from the tokenizer
212
213 Returns:
214 :obj:`str`: The HTML span for a set of consecutive chars
215 """
216 first = consecutive_chars_list[0]
217 if first.char_ix is None:
218 # its a special token
219 stoken = encoding.tokens[first.token_ix]
220 # special tokens are represented as empty spans. We use the data attribute and css
221 # magic to display it
222 return f'<span class="special-token" data-stoken={stoken}></span>'
223 # We're not in a special token so this group has a start and end.
224 last = consecutive_chars_list[-1]
225 assert first.char_ix is not None
226 assert last.char_ix is not None
227 start = first.char_ix
228 end = last.char_ix + 1
229 span_text = text[start:end]
230 css_classes = [] # What css classes will we apply on the resulting span
231 data_items = {} # What data attributes will we apply on the result span
232 if first.token_ix is not None:
233 # We can either be in a token or not (e.g. in white space)
234 css_classes.append("token")
235 if first.is_multitoken:
236 css_classes.append("multi-token")
237 if first.token_ix % 2:
238 # We use this to color alternating tokens.
239 # A token might be split by an annotation that ends in the middle of it, so this
240 # lets us visually indicate a consecutive token despite its possible splitting in
241 # the html markup
242 css_classes.append("odd-token")
243 else:
244 # Like above, but a different color so we can see the tokens alternate
245 css_classes.append("even-token")
246 if EncodingVisualizer.unk_token_regex.search(encoding.tokens[first.token_ix]) is not None:
247 # This is a special token that is in the text. probably UNK
248 css_classes.append("special-token")
249 # TODO is this the right name for the data attribute ?
250 data_items["stok"] = encoding.tokens[first.token_ix]
251 else:
252 # In this case we are looking at a group/single char that is not tokenized.
253 # e.g. white space
254 css_classes.append("non-token")
255 css = f'''class="{" ".join(css_classes)}"'''
256 data = ""
257 for key, val in data_items.items():
258 data += f' data-{key}="{val}"'
259 span_text = html.escape(span_text)
260 return f"<span {css} {data} >{span_text}</span>"
261
262 @staticmethod
263 def __make_html(text: str, encoding: Encoding, annotations: AnnotationList) -> str:
264 char_states = EncodingVisualizer.__make_char_states(text, encoding, annotations)
265 current_consecutive_chars = [char_states[0]]
266 prev_anno_ix = char_states[0].anno_ix
267 spans = []
268 label_colors_dict = EncodingVisualizer.calculate_label_colors(annotations)
269 cur_anno_ix = char_states[0].anno_ix
270 if cur_anno_ix is not None:
271 # If we started in an annotation make a span for it
272 anno = annotations[cur_anno_ix]
273 label = anno.label
274 color = label_colors_dict[label]
275 spans.append(f'<span class="annotation" style="color:{color}" data-label="{label}">')
276
277 for cs in char_states[1:]:
278 cur_anno_ix = cs.anno_ix
279 if cur_anno_ix != prev_anno_ix:
280 # If we've transitioned in or out of an annotation
281 spans.append(
282 # Create a span from the current consecutive characters
283 EncodingVisualizer.consecutive_chars_to_html(
284 current_consecutive_chars,
285 text=text,
286 encoding=encoding,
287 )
288 )
289 current_consecutive_chars = [cs]
290
291 if prev_anno_ix is not None:
292 # if we transitioned out of an annotation close it's span
293 spans.append("</span>")
294 if cur_anno_ix is not None:
295 # If we entered a new annotation make a span for it
296 anno = annotations[cur_anno_ix]
297 label = anno.label
298 color = label_colors_dict[label]
299 spans.append(f'<span class="annotation" style="color:{color}" data-label="{label}">')
300 prev_anno_ix = cur_anno_ix
301
302 if cs.partition_key() == current_consecutive_chars[0].partition_key():
303 # If the current charchter is in the same "group" as the previous one
304 current_consecutive_chars.append(cs)
305 else:
306 # Otherwise we make a span for the previous group
307 spans.append(
308 EncodingVisualizer.consecutive_chars_to_html(
309 current_consecutive_chars,
310 text=text,
311 encoding=encoding,
312 )
313 )
314 # An reset the consecutive_char_list to form a new group
315 current_consecutive_chars = [cs]
316 # All that's left is to fill out the final span
317 # TODO I think there is an edge case here where an annotation's span might not close
318 spans.append(
319 EncodingVisualizer.consecutive_chars_to_html(
320 current_consecutive_chars,
321 text=text,
322 encoding=encoding,
323 )
324 )
325
326 # Close any remaining open annotation span
327 if cur_anno_ix is not None:
328 spans.append("</span>")
329
330 res = HTMLBody(spans) # Send the list of spans to the body of our html
331 return res
332
333 @staticmethod
334 def __make_anno_map(text: str, annotations: AnnotationList) -> PartialIntList:
335 """
336 Args:
337 text (:obj:`str`):
338 The raw text we want to align to
339
340 annotations (:obj:`AnnotationList`):
341 A (possibly empty) list of annotations
342
343 Returns:
344 A list of length len(text) whose entry at index i is None if there is no annotation on
345 character i or k, the index of the annotation that covers index i where k is with
346 respect to the list of annotations
347 """
348 annotation_map = [None] * len(text)
349 for anno_ix, a in enumerate(annotations):
350 for i in range(a.start, a.end):
351 annotation_map[i] = anno_ix
352 return annotation_map
353
354 @staticmethod
355 def __make_char_states(text: str, encoding: Encoding, annotations: AnnotationList) -> List[CharState]:
356 """
357 For each character in the original text, we emit a tuple representing it's "state":
358
359 * which token_ix it corresponds to
360 * which word_ix it corresponds to
361 * which annotation_ix it corresponds to
362
363 Args:
364 text (:obj:`str`):
365 The raw text we want to align to
366
367 annotations (:obj:`List[Annotation]`):
368 A (possibly empty) list of annotations
369
370 encoding: (:class:`~tokenizers.Encoding`):
371 The encoding returned from the tokenizer
372
373 Returns:
374 :obj:`List[CharState]`: A list of CharStates, indicating for each char in the text what
375 it's state is
376 """
377 annotation_map = EncodingVisualizer.__make_anno_map(text, annotations)
378 # Todo make this a dataclass or named tuple
379 char_states: List[CharState] = [CharState(char_ix) for char_ix in range(len(text))]
380 for token_ix, token in enumerate(encoding.tokens):
381 offsets = encoding.token_to_chars(token_ix)
382 if offsets is not None:
383 start, end = offsets
384 for i in range(start, end):
385 char_states[i].tokens.append(token_ix)
386 for char_ix, anno_ix in enumerate(annotation_map):
387 char_states[char_ix].anno_ix = anno_ix
388
389 return char_states
390
391
392def HTMLBody(children: List[str], css_styles=css) -> str:
393 """
394 Generates the full html with css from a list of html spans
395
396 Args:
397 children (:obj:`List[str]`):
398 A list of strings, assumed to be html elements
399
400 css_styles (:obj:`str`, `optional`):
401 Optional alternative implementation of the css
402
403 Returns:
404 :obj:`str`: An HTML string with style markup
405 """
406 children_text = "".join(children)
407 return f"""
408 <html>
409 <head>
410 <style>
411 {css_styles}
412 </style>
413 </head>
414 <body>
415 <div class="tokenized-text" dir=auto>
416 {children_text}
417 </div>
418 </body>
419 </html>
420 """
421