ctok

#!/usr/bin/env python

import sys
import fire
from itertools import cycle
from transformers import AutoTokenizer

def cycle_colors(tokenizer_name):
    text = sys.stdin.read()
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    tokens = tokenizer.tokenize(text)

    colors = cycle(['41', '42', '43', '44', '45', '46', '47'])
    total_chars = 0

    for token in tokens:
        token_str = tokenizer.convert_tokens_to_string([token])
        token_str = token_str.replace('\n', '¶\033[0m\n')
        total_chars += len(token_str)
        print(f'\033[30;{next(colors)}m{token_str}\033[0m', end='')

    print(f"\033[0m\033[J\n{len(tokens)} tokens, {total_chars / len(tokens)} chars/token\033[0m")

def main():
    fire.Fire(cycle_colors)

if __name__ == "__main__":
    main()