from pathlib import Path import aiofiles from transformers import GPT2TokenizerFast class UsageService: def __init__(self, data_dir: Path): self.usage_file_path = data_dir / "usage.txt" # If the usage.txt file doesn't currently exist in the directory, create it and write 0.00 to it. if not self.usage_file_path.exists(): with self.usage_file_path.open("w") as f: f.write("0.00") f.close() self.tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") async def update_usage(self, tokens_used, embeddings=False): tokens_used = int(tokens_used) if not embeddings: price = ( tokens_used / 1000 ) * 0.02 # Just use the highest rate instead of model-based... I am overestimating on purpose. else: price = (tokens_used / 1000) * 0.0004 usage = await self.get_usage() print( f"Cost -> Old: {str(usage)} | New: {str(usage + float(price))}, used {str(float(price))} credits" ) # Do the same as above but with aiofiles async with aiofiles.open(self.usage_file_path, "w") as f: await f.write(str(usage + float(price))) await f.close() async def set_usage(self, usage): async with aiofiles.open(self.usage_file_path, "w") as f: await f.write(str(usage)) await f.close() async def get_usage(self): async with aiofiles.open(self.usage_file_path, "r") as f: usage = float((await f.read()).strip()) await f.close() return usage def count_tokens(self, text): res = self.tokenizer(text)["input_ids"] return len(res) async def update_usage_image(self, image_size): # 1024×1024 $0.020 / image # 512×512 $0.018 / image # 256×256 $0.016 / image if image_size == "1024x1024": price = 0.02 elif image_size == "512x512": price = 0.018 elif image_size == "256x256": price = 0.016 else: raise ValueError("Invalid image size") usage = await self.get_usage() async with aiofiles.open(self.usage_file_path, "w") as f: await f.write(str(usage + float(price))) await f.close() @staticmethod def count_tokens_static(text): tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") res = tokenizer(text)["input_ids"] return len(res)