DazedTL/modules/atelier.py
2024-01-20 19:06:25 -06:00

384 lines
No EOL
13 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import os
from pathlib import Path
import re
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
# Open AI
load_dotenv()
if os.getenv('api').replace(' ', '') != '':
openai.api_base = os.getenv('api')
openai.organization = os.getenv('org')
openai.api_key = os.getenv('key')
#Globals
MODEL = os.getenv('model')
TIMEOUT = int(os.getenv('timeout'))
LANGUAGE=os.getenv('language').capitalize()
INPUTAPICOST = .002 # Depends on the model https://openai.com/pricing
OUTPUTAPICOST = .002
PROMPT = Path('prompt.txt').read_text(encoding='utf-8')
VOCAB = Path('vocab.txt').read_text(encoding='utf-8')
THREADS = int(os.getenv('threads')) # Controls how many threads are working on a single file (May have to drop this)
LOCK = threading.Lock()
WIDTH = int(os.getenv('width'))
LISTWIDTH = int(os.getenv('listWidth'))
NOTEWIDTH = 40
MAXHISTORY = 10
ESTIMATE = ''
totalTokens = [0, 0]
NAMESLIST = []
#tqdm Globals
BAR_FORMAT='{l_bar}{bar:10}{r_bar}{bar:-10b}'
POSITION=0
LEAVE=False
# Translation Flags
FIXTEXTWRAP = True
IGNORETLTEXT = True
def handleAtelier(filename, estimate):
global ESTIMATE, totalTokens
ESTIMATE = estimate
if estimate:
start = time.time()
translatedData = openFiles(filename)
# Print Result
end = time.time()
tqdm.write(getResultString(translatedData, end - start, filename))
with LOCK:
totalTokens[0] += translatedData[1][0]
totalTokens[1] += translatedData[1][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)
outFile.writelines(translatedData[0])
# Print Result
end = time.time()
tqdm.write(getResultString(translatedData, end - start, filename))
with LOCK:
totalTokens[0] += translatedData[1][0]
totalTokens[1] += translatedData[1][1]
except Exception:
return 'Fail'
return getResultString(['', totalTokens, None], end - start, 'TOTAL')
def openFiles(filename):
with open('files/' + filename, 'r', encoding='UTF-8') as f:
translatedData = parseText(f, filename)
return translatedData
def getResultString(translatedData, translationTime, filename):
# File Print String
totalTokenstring =\
Fore.YELLOW +\
'[Input: ' + str(translatedData[1][0]) + ']'\
'[Output: ' + str(translatedData[1][1]) + ']'\
'[Cost: ${:,.4f}'.format((translatedData[1][0] * .001 * INPUTAPICOST) +\
(translatedData[1][1] * .001 * OUTPUTAPICOST)) + ']'
timeString = Fore.BLUE + '[' + str(round(translationTime, 1)) + 's]'
if translatedData[2] is None:
# Success
return filename + ': ' + totalTokenstring + timeString + Fore.GREEN + u' \u2713 ' + Fore.RESET
else:
# Fail
try:
raise translatedData[2]
except Exception as e:
errorString = str(e) + Fore.RED
return filename + ': ' + totalTokenstring + timeString + Fore.RED + u' \u2717 ' +\
errorString + Fore.RESET
def parseText(data, filename):
totalLines = 0
global LOCK
# Get total for progress bar
linesList = data.readlines()
totalLines = len(linesList)
with tqdm(bar_format=BAR_FORMAT, position=POSITION, total=totalLines, leave=LEAVE) as pbar:
pbar.desc=filename
pbar.total=totalLines
try:
response = translateText(linesList, pbar)
except Exception as e:
traceback.print_exc()
return [linesList, 0, e]
return [response[0], response[1], None]
def translateText(data, pbar):
textHistory = []
maxHistory = MAXHISTORY
totalTokens = [0,0]
syncIndex = 0
for i in range(len(data)):
if syncIndex > i:
i = syncIndex
match = re.findall(r'◆.+◆(.+)', data[i])
if len(match) > 0:
jaString = match[0]
### Translate
# Remove any textwrap
finalJAString = re.sub(r'\\n', ' ', jaString)
# Translate
response = translateGPT(finalJAString, 'Previous Text for Context: ' + ' '.join(textHistory), True)
totalTokens[0] += response[1][0]
totalTokens[1] += response[1][1]
translatedText = response[0]
# TextHistory is what we use to give GPT Context, so thats appended here.
textHistory.append('\"' + translatedText + '\"')
# Keep textHistory list at length maxHistory
if len(textHistory) > maxHistory:
textHistory.pop(0)
# Textwrap
translatedText = textwrap.fill(translatedText, width=WIDTH)
translatedText = translatedText.replace('\n', '\\n')
# Write
data[i] = data[i].replace(match[0], translatedText)
syncIndex = i + 1
pbar.update()
return [data, totalTokens]
def subVars(jaString):
jaString = jaString.replace('\u3000', ' ')
# Nested
count = 0
nestedList = re.findall(r'[\\]+[\w]+\[[\\]+[\w]+\[[0-9]+\]\]', jaString)
nestedList = set(nestedList)
if len(nestedList) != 0:
for icon in nestedList:
jaString = jaString.replace(icon, '{Nested_' + str(count) + '}')
count += 1
# Icons
count = 0
iconList = re.findall(r'[\\]+[iIkKwWaA]+\[[0-9]+\]', jaString)
iconList = set(iconList)
if len(iconList) != 0:
for icon in iconList:
jaString = jaString.replace(icon, '{Ascii_' + 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, '{Color_' + str(count) + '}')
count += 1
# Names
count = 0
nameList = re.findall(r'[\\]+[nN]\[.+?\]+', 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, '{Var_' + str(count) + '}')
count += 1
# Formatting
count = 0
if '笑えるよね.' in jaString:
print('t')
formatList = re.findall(r'[\\]+[\w]+\[.+?\]', jaString)
formatList = set(formatList)
if len(formatList) != 0:
for var in formatList:
jaString = jaString.replace(var, '{FCode_' + str(count) + '}')
count += 1
# Put all lists in list and return
allList = [nestedList, 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.strip()
translatedText = translatedText.replace(match, text)
# Nested
count = 0
if len(allList[0]) != 0:
for var in allList[0]:
translatedText = translatedText.replace('{Nested_' + str(count) + '}', var)
count += 1
# Icons
count = 0
if len(allList[1]) != 0:
for var in allList[1]:
translatedText = translatedText.replace('{Ascii_' + str(count) + '}', var)
count += 1
# Colors
count = 0
if len(allList[2]) != 0:
for var in allList[2]:
translatedText = translatedText.replace('{Color_' + str(count) + '}', var)
count += 1
# Names
count = 0
if len(allList[3]) != 0:
for var in allList[3]:
translatedText = translatedText.replace('{N_' + str(count) + '}', var)
count += 1
# Vars
count = 0
if len(allList[4]) != 0:
for var in allList[4]:
translatedText = translatedText.replace('{Var_' + str(count) + '}', var)
count += 1
# Formatting
count = 0
if len(allList[5]) != 0:
for var in allList[5]:
translatedText = translatedText.replace('{FCode_' + 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):
# Sub Vars
varResponse = subVars(t)
subbedT = varResponse[0]
# If there isn't any Japanese in the text just skip
if not re.search(r'[一-龠]+|[ぁ-ゔ]+|[ァ-ヴ]+|[\uFF00-\uFFEF]', subbedT):
return(t, [0,0])
# If ESTIMATE is True just count this as an execution and return.
if ESTIMATE:
enc = tiktoken.encoding_for_model(MODEL)
historyRaw = ''
if isinstance(history, list):
for line in history:
historyRaw += line
else:
historyRaw = history
inputTotalTokens = len(enc.encode(historyRaw)) + len(enc.encode(PROMPT))
outputTotalTokens = len(enc.encode(t)) * 2 # Estimating 2x the size of the original text
totalTokens = [inputTotalTokens, outputTotalTokens]
return (t, totalTokens)
# Characters
context = 'Game Characters:\
Character: Surname:久高 Name:有史 == Surname:Kudaka Name:Yuushi - Gender: Male\
Character: Surname:葛城 Name:碧璃 == Surname:Katsuragi Name:Midori - Gender: Female\
Character: Surname:葛城 Name:依理子 == Surname:Katsuragi Name:Yoriko - Gender: Female\
Character: Surname:桐乃木 Name:奏 == Surname:Kirinogi Name:Kanade - Gender: Female\
Character: Surname:葛城 Name:光男 == Surname:Katsuragi Name:Mitsuo - Gender: Male\
Character: Surname:尾木 Name:優真 == Surname:Ogi Name:Yuuma - Gender: Male'
# Prompt
if fullPromptFlag:
system = PROMPT
user = 'Line to Translate = ' + subbedT
else:
system = 'Output ONLY the '+ LANGUAGE +' translation in the following format: `Translation: <'+ LANGUAGE.upper() +'_TRANSLATION>`'
user = 'Line to Translate = ' + subbedT
# Create Message List
msg = []
msg.append({"role": "system", "content": system})
msg.append({"role": "user", "content": context})
if isinstance(history, list):
for line in history:
msg.append({"role": "user", "content": line})
else:
msg.append({"role": "user", "content": history})
msg.append({"role": "user", "content": user})
response = openai.ChatCompletion.create(
temperature=0,
frequency_penalty=0.2,
presence_penalty=0.2,
model=MODEL,
messages=msg,
request_timeout=TIMEOUT,
)
# Save Translated Text
translatedText = response.choices[0].message.content
totalTokens = [response.usage.prompt_tokens, response.usage.completion_tokens]
# Resub Vars
translatedText = resubVars(translatedText, varResponse[1])
# Remove Placeholder Text
translatedText = translatedText.replace(LANGUAGE +' Translation: ', '')
translatedText = translatedText.replace('Translation: ', '')
translatedText = translatedText.replace('Line to Translate = ', '')
translatedText = translatedText.replace('Translation = ', '')
translatedText = translatedText.replace('Translate = ', '')
translatedText = translatedText.replace(LANGUAGE +' Translation:', '')
translatedText = translatedText.replace('Translation:', '')
translatedText = translatedText.replace('Line to Translate =', '')
translatedText = translatedText.replace('Translation =', '')
translatedText = translatedText.replace('Translate =', '')
translatedText = translatedText.replace('', '')
translatedText = translatedText.replace('', '')
translatedText = translatedText.replace('', '')
translatedText = translatedText.replace('', '.')
translatedText = translatedText.replace('', ',')
translatedText = translatedText.replace('', '?')
translatedText = translatedText.replace('', '!')
# Return Translation
if len(translatedText) > 15 * len(t) or "I'm sorry, but I'm unable to assist with that translation" in translatedText:
raise Exception
else:
return [translatedText, totalTokens]