Count tokens for one or more files

counttokens.py

#!/usr/bin/env python3
# /// script
# dependencies = [
#   "tiktoken",
# ]
# ///
import sys
import tiktoken


def count_tokens(filepath, model="gpt-4"):
    """Count tokens in a file using the specified model's encoding."""
    encoding = tiktoken.encoding_for_model(model)

    with open(filepath, 'r') as f:
        text = f.read()

    tokens = encoding.encode(text)
    return len(tokens)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python counttokens.py <file> [file]", file=sys.stderr)
        sys.exit(1)

    total = 0
    for filepath in sys.argv[1:]:

        try:
            token_count = count_tokens(filepath)
            total += token_count
        except FileNotFoundError:
            print(f"Error: File '{filepath}' not found, skipping", file=sys.stderr)
            continue
        except Exception as e:
            print(f"Error: {e}, skipping file", file=sys.stderr)
            continue

    print(f"\nTotal token count: {total}")