Spaces:
Running on Zero
Running on Zero
| import math | |
| import unittest | |
| from html.parser import HTMLParser | |
| from highlight import describe, percent, token_cost, text_spans, render_highlight | |
| class SpanText(HTMLParser): | |
| """Collect the rendered token text: token spans plus the spaces kept between them.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.depth = 0 | |
| self.text = '' | |
| def handle_starttag(self, tag, attrs): | |
| if self.depth or dict(attrs).get('class') == 'token-text': | |
| self.depth += 1 | |
| def handle_endtag(self, tag): | |
| if self.depth: | |
| self.depth -= 1 | |
| def handle_data(self, data): | |
| if self.depth: | |
| self.text += data | |
| class HighlightTests(unittest.TestCase): | |
| def test_exact_coding_probability(self): | |
| self.assertEqual(token_cost([0, 1, 4], 0), {'bits': 2.0, 'probability': .25}) | |
| self.assertAlmostEqual(token_cost([0, 1, 4], 1)['bits'], -math.log2(.75)) | |
| def test_byte_split_unicode_whitespace_and_html(self): | |
| text = ' Hi 🌍\n你好\t<img src=x onerror="bad()"> & café ' | |
| pieces = [bytes([b]) for b in text.encode()] | |
| costs = [{'bits': 2.0, 'probability': .25}] * len(pieces) | |
| spans = text_spans(text, list(range(len(pieces))), costs, pieces) | |
| self.assertEqual(''.join(s['text'] for s in spans), text) | |
| globe = next(s for s in spans if s['text'] == '🌍') | |
| self.assertEqual(len(globe['tokens']), 4) | |
| rendered = render_highlight({'token_spans': spans}) | |
| self.assertNotIn('<img', rendered) | |
| parser = SpanText() | |
| parser.feed(rendered) | |
| self.assertEqual(parser.text, text) | |
| def test_token_details(self): | |
| self.assertEqual(describe([{'bits': 1.0, 'probability': .5}]), '1.00 bits · p 50%') | |
| self.assertEqual(describe([{'bits': 2.0, 'probability': .25}] * 3), '6.00 bits · 3 byte tokens') | |
| self.assertEqual([percent(p) for p in [.9999, .046, .0001234, 1e-7]], ['>99.9%', '4.6%', '0.012%', '<0.01%']) | |
| spans = [{'text': 'a', 'tokens': [{'bits': 1.0, 'probability': .5}]}, | |
| {'text': '\n\n', 'tokens': [{'bits': 3.0, 'probability': .125}]}, | |
| {'text': ' word', 'tokens': [{'bits': 3.0, 'probability': .125}]}, | |
| {'text': ' ', 'tokens': [{'bits': 3.0, 'probability': .125}]}] | |
| rendered = render_highlight({'token_spans': spans}) | |
| self.assertEqual(rendered.count('data-newline'), 1) | |
| self.assertIn('\n\n</span> <span', rendered) | |
| self.assertIn('data-space="2"', rendered) | |
| self.assertIn('>word</span><span', rendered) | |
| self.assertIn('> </span>', rendered) | |
| def test_empty_and_mismatched_text(self): | |
| self.assertEqual(text_spans('', [], [], []), []) | |
| self.assertIn('no tokens', render_highlight({'token_spans': []})) | |
| with self.assertRaises(ValueError): | |
| text_spans('b', [1], [{'bits': 1}], [b'a']) | |
| if __name__ == '__main__': | |
| unittest.main() | |