DazedTL/modules/json.py
2023-10-12 10:51:47 -05:00

321 lines
No EOL
10 KiB
Python

from concurrent.futures import ThreadPoolExecutor, as_completed
import json
import os
from pathlib import Path
import re
import sys
import textwrap
import threading
import time
import traceback
import tiktoken
from colorama import Fore
from dotenv import load_dotenv
import openai
from retry import retry
from tqdm import tqdm
#Globals
load_dotenv()
openai.organization = os.getenv('org')
openai.api_key = os.getenv('key')
APICOST = .002 # Depends on the model https://openai.com/pricing
PROMPT = Path('prompt.txt').read_text(encoding='utf-8')
THREADS = 10 # For GPT4 rate limit will be hit if you have more than 1 thread.
LOCK = threading.Lock()
WIDTH = 90
LISTWIDTH = 60
MAXHISTORY = 10
ESTIMATE = ''
TOTALCOST = 0
TOTALTOKENS = 0
NAMESLIST = []
#tqdm Globals
BAR_FORMAT='{l_bar}{bar:10}{r_bar}{bar:-10b}'
POSITION=0
LEAVE=False
BRFLAG = False # If the game uses <br> instead
FIXTEXTWRAP = True
def handleJSON(filename, estimate):
global ESTIMATE, TOKENS, TOTALTOKENS, TOTALCOST
ESTIMATE = estimate
if estimate:
start = time.time()
translatedData = openFiles(filename)
# Print Result
end = time.time()
tqdm.write(getResultString(translatedData, end - start, filename))
with LOCK:
TOTALCOST += translatedData[1] * .001 * APICOST
TOTALTOKENS += translatedData[1]
return getResultString(['', TOTALTOKENS, None], end - start, 'TOTAL')
else:
try:
with open('translated/' + filename, 'w', encoding='UTF-8') as outFile:
start = time.time()
translatedData = openFiles(filename)
# Print Result
end = time.time()
json.dump(translatedData[0], outFile, ensure_ascii=False)
tqdm.write(getResultString(translatedData, end - start, filename))
with LOCK:
TOTALCOST += translatedData[1] * .001 * APICOST
TOTALTOKENS += translatedData[1]
except Exception as e:
traceback.print_exc()
return 'Fail'
return getResultString(['', TOTALTOKENS, None], end - start, 'TOTAL')
def openFiles(filename):
with open('files/' + filename, 'r', encoding='UTF-8-sig') as f:
data = json.load(f)
# Map Files
if 'script' in filename:
translatedData = parseJSON(data, filename)
else:
raise NameError(filename + ' Not Supported')
return translatedData
def getResultString(translatedData, translationTime, filename):
# File Print String
tokenString = Fore.YELLOW + '[' + str(translatedData[1]) + \
' Tokens/${:,.4f}'.format(translatedData[1] * .001 * APICOST) + ']'
timeString = Fore.BLUE + '[' + str(round(translationTime, 1)) + 's]'
if translatedData[2] == None:
# Success
return filename + ': ' + tokenString + timeString + Fore.GREEN + u' \u2713 ' + Fore.RESET
else:
# Fail
try:
raise translatedData[2]
except Exception as e:
errorString = str(e) + Fore.RED
return filename + ': ' + tokenString + timeString + Fore.RED + u' \u2717 ' +\
errorString + Fore.RESET
def parseJSON(data, filename):
totalTokens = 0
totalLines = 0
totalLines = len(data)
global LOCK
with tqdm(bar_format=BAR_FORMAT, position=POSITION, total=totalLines, leave=LEAVE) as pbar:
pbar.desc=filename
pbar.total=totalLines
try:
totalTokens += translateJSON(data, pbar)
except Exception as e:
return [data, totalTokens, e]
return [data, totalTokens, None]
def translateJSON(data, pbar):
textHistory = []
maxHistory = MAXHISTORY
tokens = 0
for key, value in data.items():
# Remove any textwrap
if FIXTEXTWRAP == True:
value = re.sub(r'@b', ' ', value)
# Translate
if value == '':
response = translateGPT(key, 'Past Translated Text: ' + '|\n\n'.join(textHistory), True)
tokens += response[1]
translatedText = response[0]
textHistory.append('\"' + translatedText + '\"')
else:
translatedText = value
textHistory.append('\"' + translatedText + '\"')
# Textwrap
translatedText = textwrap.fill(translatedText, width=WIDTH)
translatedText = translatedText.replace('\n', '@b')
# Set Data
data[key] = translatedText
# Keep textHistory list at length maxHistory
if len(textHistory) > maxHistory:
textHistory.pop(0)
currentGroup = []
pbar.update(1)
return tokens
def subVars(jaString):
jaString = jaString.replace('\u3000', ' ')
# Icons
count = 0
iconList = re.findall(r'[\\]+[iI]\[[0-9]+\]', jaString)
iconList = set(iconList)
if len(iconList) != 0:
for icon in iconList:
jaString = jaString.replace(icon, '<I' + str(count) + '>')
count += 1
# Colors
count = 0
colorList = re.findall(r'[\\]+[cC]\[[0-9]+\]', jaString)
colorList = set(colorList)
if len(colorList) != 0:
for color in colorList:
jaString = jaString.replace(color, '<C' + str(count) + '>')
count += 1
# Names
count = 0
nameList = re.findall(r'[\\]+[nN]\[[0-9]+\]', jaString)
nameList = set(nameList)
if len(nameList) != 0:
for name in nameList:
jaString = jaString.replace(name, '<N' + str(count) + '>')
count += 1
# Variables
count = 0
varList = re.findall(r'[\\]+[vV]\[[0-9]+\]', jaString)
varList = set(varList)
if len(varList) != 0:
for var in varList:
jaString = jaString.replace(var, '<V' + str(count) + '>')
count += 1
# Formatting
count = 0
formatList = re.findall(r'[\\]+[!.]', jaString)
formatList = set(formatList)
if len(formatList) != 0:
for format in formatList:
jaString = jaString.replace(format, '<F' + str(count) + '>')
count += 1
# Put all lists in list and return
allList = [iconList, colorList, nameList, varList, formatList]
return [jaString, allList]
def resubVars(translatedText, allList):
# Fix Spacing and ChatGPT Nonsense
matchList = re.findall(r'<\s?.+?\s?>', translatedText)
if len(matchList) > 0:
for match in matchList:
text = match.replace(' ', '')
translatedText = translatedText.replace(match, text)
# Icons
count = 0
if len(allList[0]) != 0:
for var in allList[0]:
translatedText = translatedText.replace('<I' + str(count) + '>', var)
count += 1
# Colors
count = 0
if len(allList[1]) != 0:
for var in allList[1]:
translatedText = translatedText.replace('<C' + str(count) + '>', var)
count += 1
# Names
count = 0
if len(allList[2]) != 0:
for var in allList[2]:
translatedText = translatedText.replace('<N' + str(count) + '>', var)
count += 1
# Vars
count = 0
if len(allList[3]) != 0:
for var in allList[3]:
translatedText = translatedText.replace('<V' + str(count) + '>', var)
count += 1
# Formatting
count = 0
if len(allList[4]) != 0:
for var in allList[4]:
translatedText = translatedText.replace('<F' + str(count) + '>', var)
count += 1
# Remove Color Variables Spaces
# if '\\c' in translatedText:
# translatedText = re.sub(r'\s*(\\+c\[[1-9]+\])\s*', r' \1', translatedText)
# translatedText = re.sub(r'\s*(\\+c\[0+\])', r'\1', translatedText)
return translatedText
@retry(exceptions=Exception, tries=5, delay=5)
def translateGPT(t, history, fullPromptFlag):
# If ESTIMATE is True just count this as an execution and return.
if ESTIMATE:
enc = tiktoken.encoding_for_model("gpt-3.5-turbo")
tokens = len(enc.encode(t)) * 2 + len(enc.encode(history)) + len(enc.encode(PROMPT))
return (t, tokens)
# Sub Vars
varResponse = subVars(t)
subbedT = varResponse[0]
# If there isn't any Japanese in the text just skip
if not re.search(r'[一-龠]+|[ぁ-ゔ]+|[ァ-ヴー]+', subbedT):
return(t, 0)
"""Translate text using GPT"""
context = 'Eroge Names Context: Name: 稲盛 楓 == Inamori Kaede\nNicknames: かえちゃん == Kae-chan or かえねぇ == Kae-nee\nGender: Female,\nName: 稲盛 真守 == Inamori Mamoru\nNicknames: まーくん == Maa-kun\nGender: Male,\nName: 蓮見 雄次郎 == Hasumi Yujiro\nGender: Male,\nName: 桐谷 拓馬 == Kiriya Takuma\nNicknames: たっくん == Tak-kun\nGender: Male,\nName: 稲盛 美玖 == Inamori Yuki\nGender: Female,\nName: 奥さん == Missus\nGender: Female'
if fullPromptFlag:
system = PROMPT
user = 'Line to Translate: ' + subbedT
else:
system = 'You are an expert translator who translates everything to English. Reply with only the English Translation of the text.'
user = 'Line to Translate: ' + subbedT
response = openai.ChatCompletion.create(
temperature=0,
frequency_penalty=0.2,
presence_penalty=0.2,
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": system},
{"role": "user", "content": context},
{"role": "user", "content": history},
{"role": "user", "content": user}
],
request_timeout=30,
)
# Save Translated Text
translatedText = response.choices[0].message.content
tokens = response.usage.total_tokens
# Resub Vars
translatedText = resubVars(translatedText, varResponse[1])
# Remove Placeholder Text
translatedText = translatedText.replace('English Translation: ', '')
translatedText = translatedText.replace('Translation: ', '')
translatedText = translatedText.replace('Line to Translate: ', '')
translatedText = translatedText.replace('English Translation:', '')
translatedText = translatedText.replace('Translation:', '')
translatedText = translatedText.replace('Line to Translate:', '')
translatedText = re.sub(r'\n\nPast Translated Text:.*', '', translatedText, 0, re.DOTALL)
translatedText = re.sub(r'Note:.*', '', translatedText)
# Return Translation
if len(translatedText) > 15 * len(t) or "I'm sorry, but I'm unable to assist with that translation" in translatedText:
return [t, response.usage.total_tokens]
else:
return [translatedText, tokens]