560 lines
18 KiB
Python
560 lines
18 KiB
Python
# Libraries
|
||
import json
|
||
import os
|
||
import re
|
||
import util.dazedwrap as dazedwrap
|
||
import threading
|
||
import time
|
||
import traceback
|
||
import openai
|
||
from pathlib import Path
|
||
from colorama import Fore
|
||
from dotenv import load_dotenv
|
||
from retry import retry
|
||
from tqdm import tqdm
|
||
from util.translation import TranslationConfig, translateAI as sharedtranslateAI, getPricingConfig, calculateCost
|
||
|
||
# Open AI
|
||
load_dotenv()
|
||
if os.getenv("api").replace(" ", "") != "":
|
||
openai.base_url = 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()
|
||
PROMPT = Path("prompt.txt").read_text(encoding="utf-8")
|
||
VOCAB = Path("vocab.txt").read_text(encoding="utf-8")
|
||
THREADS = int(os.getenv("threads"))
|
||
LOCK = threading.Lock()
|
||
WIDTH = int(os.getenv("width"))
|
||
LISTWIDTH = int(os.getenv("listWidth"))
|
||
NOTEWIDTH = int(os.getenv("noteWidth"))
|
||
MAXHISTORY = 10
|
||
ESTIMATE = ""
|
||
TOKENS = [0, 0]
|
||
MISMATCH = [] # Lists files that throw a mismatch error (Length of GPT list response is wrong)
|
||
NAMESLIST = [] # List of speaker names and their translations
|
||
PBAR = None
|
||
FILENAME = None
|
||
TIMETOTAL = 0 # Total Time Taken for all translations
|
||
|
||
# Regex - Need to change this if you want to translate from/to other languages. Default is Japanese Regex
|
||
LANGREGEX = r"[\u3000-\u303F\u3040-\u309F\u30A0-\u30FF\u31F0-\u31FF\u3400-\u4DBF\u4E00-\u9FFF\uF900-\uFAFF\uFF00-\uFF5D\uFF5F-\uFFEF]+"
|
||
|
||
# Get pricing configuration based on the model
|
||
PRICING_CONFIG = getPricingConfig(MODEL)
|
||
INPUTAPICOST = PRICING_CONFIG["inputAPICost"]
|
||
OUTPUTAPICOST = PRICING_CONFIG["outputAPICost"]
|
||
BATCHSIZE = PRICING_CONFIG["batchSize"]
|
||
FREQUENCY_PENALTY = PRICING_CONFIG["frequencyPenalty"]
|
||
|
||
# tqdm Globals
|
||
BAR_FORMAT = "{l_bar}{bar:10}{r_bar}{bar:-10b}"
|
||
POSITION = 0
|
||
LEAVE = False
|
||
|
||
# Initialize Translation Config
|
||
TRANSLATION_CONFIG = TranslationConfig(
|
||
model=MODEL,
|
||
language=LANGUAGE,
|
||
prompt=PROMPT,
|
||
vocab=VOCAB,
|
||
langRegex=LANGREGEX,
|
||
batchSize=BATCHSIZE,
|
||
maxHistory=MAXHISTORY,
|
||
estimateMode=False # Will be set dynamically based on ESTIMATE
|
||
)
|
||
|
||
# Config (Default)
|
||
FIXTEXTWRAP = True # Rewrap text to WIDTH
|
||
IGNORETLTEXT = False # Skip Translated Text
|
||
|
||
# List of file patterns that use parseGeneric
|
||
# Add more patterns here as needed
|
||
GENERIC_FILES = [
|
||
"quests",
|
||
"shops",
|
||
"bonuses",
|
||
"base",
|
||
"battleprep",
|
||
"manage",
|
||
"mapcommands",
|
||
"title",
|
||
"placeevents",
|
||
"talkevents",
|
||
"characters",
|
||
|
||
# Add more file patterns here, e.g.:
|
||
# "items",
|
||
# "skills",
|
||
# "classes",
|
||
]
|
||
|
||
|
||
def handleSRPG(filename, estimate):
|
||
"""
|
||
Main handler function for SRPG Studio files.
|
||
|
||
Args:
|
||
filename: Name of the file to translate
|
||
estimate: Boolean indicating if this is an estimate run
|
||
|
||
Returns:
|
||
String with translation results or error message
|
||
"""
|
||
global ESTIMATE, TOKENS, FILENAME, TIMETOTAL
|
||
ESTIMATE = estimate
|
||
FILENAME = filename
|
||
|
||
# Translate
|
||
start = time.time()
|
||
translatedData = openFiles(filename)
|
||
|
||
# Write output file if not in estimate mode
|
||
if not estimate:
|
||
try:
|
||
with open("translated/" + filename, "w", encoding="utf-8", newline="\n") as outFile:
|
||
json.dump(translatedData[0], outFile, ensure_ascii=False, indent=4)
|
||
except Exception:
|
||
traceback.print_exc()
|
||
return "Fail"
|
||
|
||
# Print File
|
||
end = time.time()
|
||
tqdm.write(getResultString(translatedData, end - start, filename))
|
||
with LOCK:
|
||
TOKENS[0] += translatedData[1][0]
|
||
TOKENS[1] += translatedData[1][1]
|
||
|
||
# Print Total
|
||
totalString = getResultString(["", TOKENS, None], end - start, "TOTAL")
|
||
|
||
# Print any errors
|
||
if len(MISMATCH) > 0:
|
||
return totalString + Fore.RED + f"\nMismatch Errors: {MISMATCH}" + Fore.RESET
|
||
else:
|
||
return totalString
|
||
|
||
|
||
def openFiles(filename):
|
||
"""
|
||
Opens and routes SRPG Studio files to appropriate parsing functions.
|
||
|
||
Args:
|
||
filename: Name of the file to open and parse
|
||
|
||
Returns:
|
||
Tuple of (translated data, token counts, error)
|
||
"""
|
||
with open("files/" + filename, "r", encoding="utf-8-sig") as f:
|
||
data = json.load(f)
|
||
|
||
# Check if filename matches any pattern in GENERIC_FILES
|
||
if any(pattern in filename.lower() for pattern in GENERIC_FILES):
|
||
translatedData = parseGeneric(data, filename)
|
||
|
||
# TODO: Add other SRPG Studio file types here
|
||
else:
|
||
raise NameError(filename + " Not Supported")
|
||
|
||
return translatedData
|
||
|
||
|
||
def parseGeneric(data, filename):
|
||
"""
|
||
Generic parser for SRPG Studio files with id, name, desc structure.
|
||
Handles files like quests.json.
|
||
Uses a two-pass approach: first pass collects all strings to translate,
|
||
then batch translates them, then second pass applies translations.
|
||
|
||
Args:
|
||
data: Parsed JSON data (list of objects with id, name, desc)
|
||
filename: Name of the file being parsed
|
||
|
||
Returns:
|
||
Tuple of (data, token counts, error)
|
||
"""
|
||
global PBAR
|
||
|
||
totalTokens = [0, 0]
|
||
|
||
try:
|
||
# Count work units (all translatable fields that need translation)
|
||
total_units = 0
|
||
translatable_fields = ["name", "desc", "commandName", "command"]
|
||
|
||
for entry in data:
|
||
if entry:
|
||
for field in translatable_fields:
|
||
if field in entry and entry[field]:
|
||
total_units += 1
|
||
|
||
# Handle pages array separately
|
||
if "pages" in entry and entry["pages"] and isinstance(entry["pages"], list):
|
||
for page in entry["pages"]:
|
||
if page:
|
||
total_units += 1
|
||
|
||
# Setup progress bar
|
||
with LOCK:
|
||
PBAR = tqdm(
|
||
desc=filename,
|
||
total=total_units,
|
||
bar_format=BAR_FORMAT,
|
||
position=POSITION,
|
||
leave=LEAVE,
|
||
)
|
||
|
||
# Translate the data using two-pass approach
|
||
result = translateGeneric(data, filename)
|
||
totalTokens[0] += result[0]
|
||
totalTokens[1] += result[1]
|
||
|
||
# Close progress bar
|
||
with LOCK:
|
||
if PBAR:
|
||
PBAR.close()
|
||
PBAR = None
|
||
|
||
return (data, totalTokens, None)
|
||
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
if PBAR:
|
||
with LOCK:
|
||
PBAR.close()
|
||
PBAR = None
|
||
return (data, totalTokens, e)
|
||
|
||
|
||
def translateGeneric(data, filename, translatedDataList=None):
|
||
"""
|
||
Translates generic SRPG Studio data with id, name, desc, commandName, command, pages structure.
|
||
Uses two-pass approach via recursion:
|
||
- Pass 1 (translatedDataList=None): Collect strings and batch translate
|
||
- Pass 2 (translatedDataList set): Apply translations back to the data
|
||
|
||
Args:
|
||
data: List of objects with id, name, desc, commandName, command, pages keys
|
||
filename: Name of the file being translated
|
||
translatedDataList: List containing [nameList, descList, commandNameList, commandList, pagesList]
|
||
- Pass 1: Empty lists to collect originals
|
||
- Pass 2: Filled lists with translations
|
||
|
||
Returns:
|
||
Tuple of [input tokens, output tokens]
|
||
"""
|
||
global PBAR
|
||
|
||
totalTokens = [0, 0]
|
||
|
||
# Initialize or extract lists
|
||
if translatedDataList is None:
|
||
# PASS 1: Create empty lists to collect strings
|
||
nameList = []
|
||
descList = []
|
||
commandNameList = []
|
||
commandList = []
|
||
pagesList = []
|
||
else:
|
||
# PASS 2: Use provided translated lists
|
||
nameList = translatedDataList[0]
|
||
descList = translatedDataList[1]
|
||
commandNameList = translatedDataList[2]
|
||
commandList = translatedDataList[3]
|
||
pagesList = translatedDataList[4]
|
||
|
||
# Single loop - behavior depends on which pass we're in
|
||
for entry in data:
|
||
if not entry:
|
||
continue
|
||
|
||
# Handle name field
|
||
if "name" in entry and entry["name"]:
|
||
# PASS 1: Collect original
|
||
if translatedDataList is None:
|
||
nameList.append(entry["name"])
|
||
# PASS 2: Apply translation
|
||
else:
|
||
if nameList:
|
||
entry["name"] = nameList[0]
|
||
nameList.pop(0)
|
||
|
||
# Handle desc field
|
||
if "desc" in entry and entry["desc"]:
|
||
# PASS 1: Collect original
|
||
if translatedDataList is None:
|
||
descList.append(entry["desc"])
|
||
# PASS 2: Apply translation
|
||
else:
|
||
if descList:
|
||
entry["desc"] = descList[0]
|
||
descList.pop(0)
|
||
|
||
# Handle commandName field
|
||
if "commandName" in entry and entry["commandName"]:
|
||
# PASS 1: Collect original
|
||
if translatedDataList is None:
|
||
commandNameList.append(entry["commandName"])
|
||
# PASS 2: Apply translation
|
||
else:
|
||
if commandNameList:
|
||
entry["commandName"] = commandNameList[0]
|
||
commandNameList.pop(0)
|
||
|
||
# Handle command field
|
||
if "command" in entry and entry["command"]:
|
||
# PASS 1: Collect original
|
||
if translatedDataList is None:
|
||
commandList.append(entry["command"])
|
||
# PASS 2: Apply translation
|
||
else:
|
||
if commandList:
|
||
entry["command"] = commandList[0]
|
||
commandList.pop(0)
|
||
|
||
# Handle pages field (array of strings)
|
||
if "pages" in entry and entry["pages"] and isinstance(entry["pages"], list):
|
||
for i, page in enumerate(entry["pages"]):
|
||
if page:
|
||
# PASS 1: Collect original
|
||
if translatedDataList is None:
|
||
# Nuke Wordwrap
|
||
page = page.replace("\n", " ")
|
||
pagesList.append(page)
|
||
# PASS 2: Apply translation
|
||
else:
|
||
if pagesList:
|
||
translatedText = pagesList[0]
|
||
|
||
# Wordwrap
|
||
translatedText = dazedwrap.wrapText(translatedText, width=WIDTH)
|
||
|
||
# Set Data
|
||
entry["pages"][i] = translatedText
|
||
pagesList.pop(0)
|
||
|
||
# If this was Pass 1, do the translation and recurse for Pass 2
|
||
if translatedDataList is None:
|
||
# Store original counts for mismatch checking
|
||
originalNameCount = len(nameList)
|
||
originalDescCount = len(descList)
|
||
originalCommandNameCount = len(commandNameList)
|
||
originalCommandCount = len(commandList)
|
||
originalPagesCount = len(pagesList)
|
||
|
||
# Batch translate names
|
||
if nameList:
|
||
response = translateAI(
|
||
nameList,
|
||
"Reply with only the " + LANGUAGE + " translation of the quest name.",
|
||
True
|
||
)
|
||
nameList = response[0]
|
||
totalTokens[0] += response[1][0]
|
||
totalTokens[1] += response[1][1]
|
||
|
||
# Batch translate descriptions
|
||
if descList:
|
||
response = translateAI(
|
||
descList,
|
||
"Reply with only the " + LANGUAGE + " translation of the quest description.",
|
||
True
|
||
)
|
||
descList = response[0]
|
||
totalTokens[0] += response[1][0]
|
||
totalTokens[1] += response[1][1]
|
||
|
||
# Batch translate command names
|
||
if commandNameList:
|
||
response = translateAI(
|
||
commandNameList,
|
||
"Reply with only the " + LANGUAGE + " translation of the command name.",
|
||
True
|
||
)
|
||
commandNameList = response[0]
|
||
totalTokens[0] += response[1][0]
|
||
totalTokens[1] += response[1][1]
|
||
|
||
# Batch translate commands
|
||
if commandList:
|
||
response = translateAI(
|
||
commandList,
|
||
"Reply with only the " + LANGUAGE + " translation of the command.",
|
||
True
|
||
)
|
||
commandList = response[0]
|
||
totalTokens[0] += response[1][0]
|
||
totalTokens[1] += response[1][1]
|
||
|
||
# Batch translate pages
|
||
if pagesList:
|
||
response = translateAI(
|
||
pagesList,
|
||
"Reply with only the " + LANGUAGE + " translation of the page content.",
|
||
True
|
||
)
|
||
pagesList = response[0]
|
||
totalTokens[0] += response[1][0]
|
||
totalTokens[1] += response[1][1]
|
||
|
||
# Check for mismatch errors
|
||
if len(nameList) != originalNameCount:
|
||
with LOCK:
|
||
if filename not in MISMATCH:
|
||
MISMATCH.append(filename)
|
||
|
||
if len(descList) != originalDescCount:
|
||
with LOCK:
|
||
if filename not in MISMATCH:
|
||
MISMATCH.append(filename)
|
||
|
||
if len(commandNameList) != originalCommandNameCount:
|
||
with LOCK:
|
||
if filename not in MISMATCH:
|
||
MISMATCH.append(filename)
|
||
|
||
if len(commandList) != originalCommandCount:
|
||
with LOCK:
|
||
if filename not in MISMATCH:
|
||
MISMATCH.append(filename)
|
||
|
||
if len(pagesList) != originalPagesCount:
|
||
with LOCK:
|
||
if filename not in MISMATCH:
|
||
MISMATCH.append(filename)
|
||
|
||
# PASS 2: Recursively call to apply translations
|
||
translateGeneric(data, filename, [nameList, descList, commandNameList, commandList, pagesList])
|
||
|
||
return totalTokens
|
||
|
||
|
||
def getResultString(translatedData, translationTime, filename):
|
||
"""
|
||
Formats the translation result string with token counts, cost, and time.
|
||
|
||
Args:
|
||
translatedData: Tuple of (data, tokens, error)
|
||
translationTime: Time taken for translation
|
||
filename: Name of the file
|
||
|
||
Returns:
|
||
Formatted result string
|
||
"""
|
||
global TIMETOTAL
|
||
|
||
# Calculate cost
|
||
cost = calculateCost(translatedData[1][0], translatedData[1][1], MODEL)
|
||
totalTokenstring = (
|
||
Fore.YELLOW + "[Input: " + str(translatedData[1][0]) + "]"
|
||
"[Output: "
|
||
+ str(translatedData[1][1])
|
||
+ "]" "[Cost: ${:,.4f}".format(cost)
|
||
+ "]"
|
||
)
|
||
|
||
# Format time string
|
||
if filename != "TOTAL":
|
||
timeString = Fore.BLUE + "[" + str(round(translationTime, 1)) + "s]"
|
||
TIMETOTAL += round(translationTime, 1)
|
||
else:
|
||
timeString = Fore.BLUE + "[" + str(round(TIMETOTAL, 1)) + "s]"
|
||
|
||
# Return success or failure string
|
||
if translatedData[2] is None:
|
||
# Success
|
||
return filename + ": " + totalTokenstring + timeString + Fore.GREEN + " \u2713 " + Fore.RESET
|
||
else:
|
||
# Fail
|
||
try:
|
||
raise translatedData[2]
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
return (
|
||
filename
|
||
+ ": "
|
||
+ totalTokenstring
|
||
+ timeString
|
||
+ Fore.RED
|
||
+ " \u2717 "
|
||
+ Fore.RESET
|
||
)
|
||
|
||
|
||
def getSpeaker(speaker):
|
||
"""
|
||
Translates speaker/character names with caching to avoid redundant translations.
|
||
|
||
Args:
|
||
speaker: The original speaker name to translate
|
||
|
||
Returns:
|
||
List containing [translated name, [input tokens, output tokens]]
|
||
"""
|
||
if speaker == "":
|
||
return ["", [0, 0]]
|
||
|
||
# Check if speaker has already been translated
|
||
for i in range(len(NAMESLIST)):
|
||
if speaker == NAMESLIST[i][0]:
|
||
return [NAMESLIST[i][1], [0, 0]]
|
||
|
||
# Translate and Store Speaker
|
||
response = translateAI(
|
||
speaker,
|
||
"Reply with the " + LANGUAGE + " translation of the NPC name.",
|
||
False,
|
||
)
|
||
response[0] = response[0].title()
|
||
response[0] = response[0].replace("'S", "'s")
|
||
response[0] = response[0].replace("Speaker: ", "")
|
||
|
||
# Retry if name doesn't translate for some reason
|
||
if re.search(r"([a-zA-Z？?])", response[0]) is None:
|
||
response = translateAI(
|
||
speaker,
|
||
"Reply with the " + LANGUAGE + " translation of the NPC name.",
|
||
False,
|
||
)
|
||
response[0] = response[0].title()
|
||
response[0] = response[0].replace("'S", "'s")
|
||
|
||
speakerList = [speaker, response[0]]
|
||
NAMESLIST.append(speakerList)
|
||
return response
|
||
|
||
|
||
def translateAI(text, history, fullPromptFlag):
|
||
"""
|
||
Legacy wrapper function for the new shared translation utility.
|
||
This maintains compatibility with existing code while using the new shared implementation.
|
||
|
||
Args:
|
||
text: Text to translate (can be string or list)
|
||
history: History/context for the translation
|
||
fullPromptFlag: Whether to use the full prompt with vocab
|
||
|
||
Returns:
|
||
List containing [translated text, [input tokens, output tokens]]
|
||
"""
|
||
global PBAR, MISMATCH, FILENAME
|
||
|
||
# Update config estimate mode based on global ESTIMATE
|
||
TRANSLATION_CONFIG.estimateMode = bool(ESTIMATE)
|
||
|
||
# Call the new shared translation function
|
||
return sharedtranslateAI(
|
||
text=text,
|
||
history=history,
|
||
fullPromptFlag=fullPromptFlag,
|
||
config=TRANSLATION_CONFIG,
|
||
filename=FILENAME,
|
||
pbar=PBAR,
|
||
lock=LOCK,
|
||
mismatchList=MISMATCH
|
||
)
|
||
|