umar-30 commited on
Commit
14febb2
·
1 Parent(s): 53ae001

Add missing utils folder

Browse files
Files changed (2) hide show
  1. utils/__init__.py +1 -0
  2. utils/token_counter.py +51 -0
utils/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ # Utils package
utils/token_counter.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Token Counter Utility
3
+
4
+ Provides token counting functionality using tiktoken library.
5
+ Used for managing OpenAI API token limits.
6
+
7
+ Functions:
8
+ count_tokens: Count tokens in text using specified model encoding
9
+ """
10
+
11
+ import tiktoken
12
+ from typing import Optional
13
+
14
+
15
+ def count_tokens(text: str, model: str = "gpt-3.5-turbo") -> int:
16
+ """
17
+ Count tokens in text using tiktoken for specified model.
18
+
19
+ This function uses the tiktoken library to accurately count tokens
20
+ as OpenAI's API would count them. This is essential for:
21
+ - Staying within context limits
22
+ - Estimating API costs
23
+ - Managing prompt sizes
24
+
25
+ Args:
26
+ text: Text to count tokens in
27
+ model: OpenAI model name (default: "gpt-3.5-turbo")
28
+ Supported models: gpt-3.5-turbo, gpt-4, gpt-4-turbo-preview
29
+
30
+ Returns:
31
+ Number of tokens in the text
32
+
33
+ Raises:
34
+ ValueError: If model encoding is not available
35
+
36
+ Example:
37
+ >>> count_tokens("Hello, world!")
38
+ 4
39
+ >>> count_tokens("How do I create a ROS 2 publisher?")
40
+ 10
41
+ >>> count_tokens("Long text...", model="gpt-4")
42
+ ...
43
+ """
44
+ try:
45
+ encoder = tiktoken.encoding_for_model(model)
46
+ except KeyError:
47
+ # Fallback to cl100k_base encoding (used by gpt-3.5-turbo and gpt-4)
48
+ encoder = tiktoken.get_encoding("cl100k_base")
49
+
50
+ tokens = encoder.encode(text)
51
+ return len(tokens)