|
|
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)
|
