279 lines
8.4 KiB
Python
279 lines
8.4 KiB
Python
import json
|
||
import os
|
||
import re
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import traceback
|
||
from pathlib import Path
|
||
|
||
from colorama import Fore
|
||
from tqdm import tqdm
|
||
|
||
import util.dazedwrap as dazedwrap
|
||
from util.translation import (
|
||
TranslationConfig,
|
||
calculateCost,
|
||
getPricingConfig,
|
||
translateAI as sharedtranslateAI,
|
||
)
|
||
|
||
|
||
MODEL = os.getenv("model")
|
||
LANGUAGE = os.getenv("language").capitalize()
|
||
PROMPT = Path("prompt.txt").read_text(encoding="utf-8")
|
||
VOCAB = Path("vocab.txt").read_text(encoding="utf-8")
|
||
LOCK = threading.Lock()
|
||
WIDTH = int(os.getenv("width"))
|
||
MAXHISTORY = 10
|
||
ESTIMATE = ""
|
||
TOKENS = [0, 0]
|
||
NAMESLIST = []
|
||
MISMATCH = []
|
||
FILENAME = None
|
||
PBAR = None
|
||
|
||
BAR_FORMAT = "{l_bar}{bar:10}{r_bar}{bar:-10b}"
|
||
POSITION = 0
|
||
LEAVE = False
|
||
|
||
LANGREGEX = r"[一-龠ぁ-ゔァ-ヴーａ-ｚＡ-Ｚ０-９\uFF61-\uFF9F]+"
|
||
PRICING_CONFIG = getPricingConfig(MODEL)
|
||
BATCHSIZE = PRICING_CONFIG["batchSize"]
|
||
|
||
TRANSLATION_CONFIG = TranslationConfig(
|
||
model=MODEL,
|
||
language=LANGUAGE,
|
||
prompt=PROMPT,
|
||
vocab=VOCAB,
|
||
langRegex=LANGREGEX,
|
||
batchSize=BATCHSIZE,
|
||
maxHistory=MAXHISTORY,
|
||
estimateMode=False,
|
||
)
|
||
|
||
|
||
def handleYuris(filename, estimate):
|
||
global ESTIMATE, FILENAME, TOKENS
|
||
ESTIMATE = estimate
|
||
FILENAME = filename
|
||
|
||
try:
|
||
start = time.time()
|
||
translatedData = openFiles(filename)
|
||
|
||
if not estimate:
|
||
os.makedirs(os.path.dirname(os.path.join("translated", filename)) or "translated", exist_ok=True)
|
||
with open("translated/" + filename, "w", encoding="utf-8", newline="\n") as outFile:
|
||
json.dump(translatedData[0], outFile, ensure_ascii=False, indent=4)
|
||
|
||
end = time.time()
|
||
tqdm.write(getResultString(translatedData, end - start, filename))
|
||
with LOCK:
|
||
TOKENS[0] += translatedData[1][0]
|
||
TOKENS[1] += translatedData[1][1]
|
||
except Exception:
|
||
traceback.print_exc()
|
||
return "Fail"
|
||
|
||
totalString = getResultString(["", TOKENS, None], end - start, "TOTAL")
|
||
if MISMATCH:
|
||
return totalString + Fore.RED + f"\nMismatch Errors: {MISMATCH}" + Fore.RESET
|
||
return totalString
|
||
|
||
|
||
def openFiles(filename):
|
||
with open("files/" + filename, "r", encoding="UTF-8-sig") as readFile:
|
||
data = json.load(readFile)
|
||
translatedData = parseYuris(data, filename)
|
||
return translatedData
|
||
|
||
|
||
def parseYuris(data, filename):
|
||
totalTokens = [0, 0]
|
||
global PBAR
|
||
|
||
if not isinstance(data, list):
|
||
return [data, totalTokens, TypeError(f"{filename} must be a JSON array")]
|
||
|
||
totalLines = sum(
|
||
1
|
||
for item in data
|
||
if isinstance(item, dict) and isinstance(item.get("message"), str) and item["message"].strip()
|
||
)
|
||
|
||
with tqdm(bar_format=BAR_FORMAT, position=POSITION, total=totalLines, leave=LEAVE) as pbar:
|
||
pbar.desc = filename
|
||
PBAR = pbar
|
||
try:
|
||
result = translateYuris(data, filename)
|
||
totalTokens[0] += result[0]
|
||
totalTokens[1] += result[1]
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
return [data, totalTokens, e]
|
||
|
||
return [data, totalTokens, None]
|
||
|
||
|
||
def translateYuris(data, filename):
|
||
global PBAR
|
||
tokens = [0, 0]
|
||
entries = []
|
||
stringList = []
|
||
|
||
for index, item in enumerate(data):
|
||
if not isinstance(item, dict):
|
||
continue
|
||
|
||
speaker = ""
|
||
if isinstance(item.get("name"), str) and item["name"].strip():
|
||
speakerData = getSpeaker(item["name"].strip())
|
||
speaker = speakerData[0]
|
||
tokens[0] += speakerData[1][0]
|
||
tokens[1] += speakerData[1][1]
|
||
item["name"] = speaker
|
||
|
||
message = item.get("message")
|
||
if not isinstance(message, str) or not message.strip():
|
||
continue
|
||
# if not re.search(LANGREGEX, message):
|
||
# continue
|
||
|
||
cleanMessage = message.replace("\r\n", " ").replace("\n", " ").strip()
|
||
hasSpeaker = bool(speaker)
|
||
stringList.append(f"[{speaker}]: {cleanMessage}" if hasSpeaker else cleanMessage)
|
||
entries.append((index, message, hasSpeaker))
|
||
|
||
if not stringList:
|
||
return tokens
|
||
|
||
PBAR.total = len(stringList)
|
||
PBAR.refresh()
|
||
|
||
response = translateAI(stringList, "Reply with the English Translation")
|
||
tokens[0] += response[1][0]
|
||
tokens[1] += response[1][1]
|
||
translatedList = response[0]
|
||
|
||
if len(stringList) != len(translatedList):
|
||
with LOCK:
|
||
if FILENAME not in MISMATCH:
|
||
MISMATCH.append(FILENAME)
|
||
return tokens
|
||
|
||
for (index, originalMessage, hasSpeaker), translatedText in zip(entries, translatedList):
|
||
if hasSpeaker:
|
||
translatedText = stripSpeakerPrefix(translatedText)
|
||
translatedText = translatedText.replace("\r\n", "\n")
|
||
translatedText = dazedwrap.wrapText(translatedText, width=WIDTH)
|
||
translatedText = translatedText.replace("\n", "\r\n")
|
||
data[index]["message"] = originalMessage.replace(originalMessage, translatedText)
|
||
save_progress_json(data, filename)
|
||
if PBAR is not None:
|
||
PBAR.update(1)
|
||
|
||
return tokens
|
||
|
||
|
||
def stripSpeakerPrefix(translated_text):
|
||
"""Same behavior as modules/json.translateJSON: strip [Name]: prefix from TL lines."""
|
||
match = re.search(r"(^\[.+?\]\s?[|:]\s?)", translated_text)
|
||
if match:
|
||
translated_text = translated_text.replace(match.group(1), "")
|
||
else:
|
||
cjk_m = re.match(
|
||
r"^[一-龠ぁ-ゔァ-ヴーａ-ｚＡ-Ｚ０-９\uFF61-\uFF9F]+\s*[\(（:]\s*",
|
||
translated_text,
|
||
)
|
||
if cjk_m:
|
||
translated_text = translated_text[cjk_m.end():]
|
||
translated_text = re.sub(r"\s*[)）]\s*$", "", translated_text)
|
||
return translated_text.strip()
|
||
|
||
|
||
def save_progress_json(data, filename):
|
||
try:
|
||
if ESTIMATE:
|
||
return
|
||
target = os.path.join("translated", filename)
|
||
os.makedirs(os.path.dirname(target) or "translated", exist_ok=True)
|
||
tmp_fd, tmp_path = tempfile.mkstemp(
|
||
prefix=f"{os.path.basename(filename)}.",
|
||
suffix=".tmp",
|
||
dir=os.path.dirname(target) or "translated",
|
||
)
|
||
try:
|
||
with os.fdopen(tmp_fd, "w", encoding="utf-8", newline="\n") as tmp_file:
|
||
json.dump(data, tmp_file, ensure_ascii=False, indent=4)
|
||
os.replace(tmp_path, target)
|
||
finally:
|
||
if os.path.exists(tmp_path):
|
||
try:
|
||
os.remove(tmp_path)
|
||
except OSError:
|
||
pass
|
||
except Exception:
|
||
traceback.print_exc()
|
||
|
||
|
||
def getResultString(translatedData, translationTime, filename):
|
||
cost = calculateCost(translatedData[1][0], translatedData[1][1], MODEL)
|
||
totalTokenstring = (
|
||
Fore.YELLOW
|
||
+ f"[Input: {translatedData[1][0]}]"
|
||
+ f"[Output: {translatedData[1][1]}]"
|
||
+ f"[Cost: ${cost:,.4f}]"
|
||
)
|
||
timeString = Fore.BLUE + f"[{round(translationTime, 1)}s]"
|
||
|
||
if translatedData[2] is None:
|
||
return f"{filename}: {totalTokenstring}{timeString}" + Fore.GREEN + " \u2713 " + Fore.RESET
|
||
|
||
try:
|
||
raise translatedData[2]
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
errorString = str(e) + Fore.RED
|
||
return f"{filename}: {totalTokenstring}{timeString}" + Fore.RED + " \u2717 " + errorString + Fore.RESET
|
||
|
||
|
||
def getSpeaker(speaker):
|
||
if not speaker:
|
||
return ["", [0, 0]]
|
||
|
||
for original, translated in NAMESLIST:
|
||
if speaker == original:
|
||
return [translated, [0, 0]]
|
||
|
||
response = translateAI(
|
||
speaker,
|
||
"Reply with the " + LANGUAGE + " translation of the NPC name.",
|
||
False,
|
||
)
|
||
response[0] = response[0].title().replace("'S", "'s").replace("Speaker: ", "")
|
||
|
||
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().replace("'S", "'s").replace("Speaker: ", "")
|
||
|
||
NAMESLIST.append([speaker, response[0]])
|
||
return response
|
||
|
||
|
||
def translateAI(text, history, history_ctx=None):
|
||
global PBAR, MISMATCH, FILENAME
|
||
TRANSLATION_CONFIG.estimateMode = bool(ESTIMATE)
|
||
return sharedtranslateAI(
|
||
text=text,
|
||
history=history,
|
||
config=TRANSLATION_CONFIG,
|
||
filename=FILENAME,
|
||
pbar=PBAR,
|
||
lock=LOCK,
|
||
mismatchList=MISMATCH,
|
||
)
|