chat / app /utils /token_counter.py
ariansyahdedy's picture
Add prompt edit and api key config
8d2f9d4
raw
history blame
775 Bytes
# token_counter.py
import os
import tiktoken
# Choose the encoding based on your model, e.g., 'cl100k_base' for OpenAI models
encoding = tiktoken.get_encoding("cl100k_base")
def count_tokens(text):
tokens = encoding.encode(text)
return len(tokens)
class TokenCounter:
def __init__(self):
self.total_tokens = 0
self.doc_tokens = {}
def add_document(self, doc_id, text):
num_tokens = count_tokens(text)
self.doc_tokens[doc_id] = num_tokens
self.total_tokens += num_tokens
def remove_document(self, doc_id):
if doc_id in self.doc_tokens:
self.total_tokens -= self.doc_tokens[doc_id]
del self.doc_tokens[doc_id]
def get_total_tokens(self):
return self.total_tokens