count_tokens.py 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. #!/usr/bin/env python3
  2. # /// script
  3. # requires-python = ">=3.9"
  4. # dependencies = ["tiktoken"]
  5. # ///
  6. """count_tokens — the single length metric for skill authoring.
  7. Token counts replace line counts everywhere in the builder and eval-runner.
  8. This script reports the token length of a file or of text piped on stdin, using
  9. the tiktoken cl100k_base encoding. When tiktoken is not installed it falls back
  10. to a character-based estimate (len(text) // 4) and says so, so the script always
  11. runs under a bare python3 even with no third-party packages present.
  12. Usage:
  13. count_tokens.py <file> count the tokens in a file
  14. count_tokens.py --stdin count the tokens read from stdin
  15. Output (one line of JSON on stdout):
  16. {"tokens": <int>, "method": "tiktoken"} when tiktoken loaded
  17. {"tokens": <int>, "method": "fallback"} when it fell back to chars // 4
  18. Budgets this feeds: SKILL.md ~1500-2500, multi-branch reference ~4500,
  19. single-purpose reference ~9000.
  20. """
  21. import argparse
  22. import json
  23. import sys
  24. ENCODING = "cl100k_base"
  25. def count_tokens(text: str) -> tuple[int, str]:
  26. """Return (token_count, method).
  27. Tries tiktoken's cl100k_base encoding first. If tiktoken cannot be imported
  28. or initialized, estimates with len(text) // 4 and reports method "fallback".
  29. """
  30. try:
  31. import tiktoken
  32. except Exception:
  33. return len(text) // 4, "fallback"
  34. try:
  35. enc = tiktoken.get_encoding(ENCODING)
  36. except Exception:
  37. return len(text) // 4, "fallback"
  38. return len(enc.encode(text)), "tiktoken"
  39. def read_input(args) -> str:
  40. if args.stdin:
  41. return sys.stdin.read()
  42. with open(args.file, encoding="utf-8") as f:
  43. return f.read()
  44. def main(argv: list[str] | None = None) -> int:
  45. p = argparse.ArgumentParser(
  46. description=__doc__,
  47. formatter_class=argparse.RawDescriptionHelpFormatter,
  48. )
  49. p.add_argument("file", nargs="?", help="path to the file to count")
  50. p.add_argument("--stdin", action="store_true", help="read text from stdin instead of a file")
  51. args = p.parse_args(argv)
  52. if not args.stdin and not args.file:
  53. p.error("provide a file path or --stdin")
  54. if args.stdin and args.file:
  55. p.error("provide either a file path or --stdin, not both")
  56. text = read_input(args)
  57. tokens, method = count_tokens(text)
  58. print(json.dumps({"tokens": tokens, "method": method}))
  59. return 0
  60. if __name__ == "__main__":
  61. sys.exit(main())