CLI utility for estimating tokens

token

#!/usr/bin/env python3.11

# Script for estimating tokens
# Installation: (MacOS)
# - Copy this gist in ~/.local/bin and call it token
# - chmod +x ~/.local/bin/token
# - pip3 install tiktoken
# - echo "hello world"|token

import sys
import tiktoken

def count_tokens(text):
    encoding = tiktoken.get_encoding("cl100k_base")
    return len(encoding.encode(text))

if __name__ == "__main__":
    if sys.stdin.isatty():  # Check if input is from a terminal
        if len(sys.argv) < 2:
            print("Usage: tokens <text> OR cat text.txt | tokens")
            sys.exit(1)
        text = " ".join(sys.argv[1:])
    else:  # Read from stdin
        text = sys.stdin.read().strip()

    token_count = count_tokens(text)
    print(f"Token count: {token_count}")