229 lines
7.6 KiB
Python
229 lines
7.6 KiB
Python
"""Translate WolfDawn extraction JSON.
|
||
|
||
This module consumes the JSON produced by the vendored WolfDawn ``wolf`` CLI
|
||
(``strings-extract`` / ``names-extract``) and fills in the ``text`` fields with
|
||
translations, leaving ``source`` untouched so WolfDawn's inject drift-guard can
|
||
match each line back to its original.
|
||
|
||
Supported document ``kind``s (all share the ``{source, text}`` leaf pattern):
|
||
map / common -> scenes[].lines[]
|
||
db -> groups[].lines[]
|
||
gamedat -> lines[]
|
||
txt -> lines[]
|
||
txt-dir -> files[].lines[]
|
||
names -> names[]
|
||
|
||
Only entries whose ``source`` contains target-language (Japanese by default)
|
||
text are sent to the model; everything else keeps ``text == source`` so inject
|
||
is a no-op for it. Translated text is written back verbatim (no re-wrapping) to
|
||
keep WolfDawn's inline-code and byte-exact guards happy.
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import threading
|
||
import time
|
||
import traceback
|
||
|
||
from colorama import Fore
|
||
from tqdm import tqdm
|
||
|
||
from util.paths import PROMPT_PATH, VOCAB_PATH
|
||
from util.translation import (
|
||
TranslationConfig,
|
||
translateAI as sharedtranslateAI,
|
||
getPricingConfig,
|
||
calculateCost,
|
||
)
|
||
|
||
# Globals (mirror the other engine modules; populated from .env at import time)
|
||
MODEL = os.getenv("model")
|
||
TIMEOUT = int(os.getenv("timeout"))
|
||
LANGUAGE = os.getenv("language").capitalize()
|
||
|
||
PROMPT = PROMPT_PATH.read_text(encoding="utf-8")
|
||
VOCAB = VOCAB_PATH.read_text(encoding="utf-8")
|
||
LOCK = threading.Lock()
|
||
MAXHISTORY = 10
|
||
ESTIMATE = ""
|
||
TOKENS = [0, 0]
|
||
MISMATCH = [] # Files that hit a length-mismatch during translation
|
||
FILENAME = None
|
||
|
||
# Regex - default matches Japanese (kanji, kana, full-width forms).
|
||
LANGREGEX = r"[一-龠ぁ-ゔァ-ヴーａ-ｚＡ-Ｚ０-９\uFF61-\uFF9F]+"
|
||
|
||
# Pricing / batching from the configured model
|
||
PRICING_CONFIG = getPricingConfig(MODEL)
|
||
INPUTAPICOST = PRICING_CONFIG["inputAPICost"]
|
||
OUTPUTAPICOST = PRICING_CONFIG["outputAPICost"]
|
||
BATCHSIZE = PRICING_CONFIG["batchSize"]
|
||
FREQUENCY_PENALTY = PRICING_CONFIG["frequencyPenalty"]
|
||
|
||
# tqdm / progress globals (PBAR is polled by util/subprocess_runner.py)
|
||
BAR_FORMAT = "{l_bar}{bar:10}{r_bar}{bar:-10b}"
|
||
POSITION = 0
|
||
LEAVE = False
|
||
PBAR = None
|
||
|
||
TRANSLATION_CONFIG = TranslationConfig(
|
||
model=MODEL,
|
||
language=LANGUAGE,
|
||
prompt=PROMPT,
|
||
vocab=VOCAB,
|
||
langRegex=LANGREGEX,
|
||
batchSize=BATCHSIZE,
|
||
maxHistory=MAXHISTORY,
|
||
estimateMode=False,
|
||
)
|
||
|
||
|
||
def handleWolfDawn(filename, estimate):
|
||
"""Entry point used by the CLI/GUI dispatchers. Returns a summary string or 'Fail'."""
|
||
global ESTIMATE, TOKENS, FILENAME
|
||
ESTIMATE = estimate
|
||
FILENAME = filename
|
||
|
||
start = time.time()
|
||
translatedData = openFiles(filename)
|
||
|
||
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"
|
||
|
||
end = time.time()
|
||
tqdm.write(getResultString(translatedData, end - start, filename))
|
||
with LOCK:
|
||
TOKENS[0] += translatedData[1][0]
|
||
TOKENS[1] += translatedData[1][1]
|
||
|
||
totalString = getResultString(["", TOKENS, None], end - start, "TOTAL")
|
||
if len(MISMATCH) > 0:
|
||
return totalString + Fore.RED + f"\nMismatch Errors: {MISMATCH}" + Fore.RESET
|
||
return totalString
|
||
|
||
|
||
def openFiles(filename):
|
||
"""Load the extraction JSON, translate it in place, and return [data, tokens, error]."""
|
||
with open("files/" + filename, "r", encoding="utf-8-sig") as f:
|
||
data = json.load(f)
|
||
|
||
kind = data.get("kind")
|
||
if kind not in ("map", "common", "db", "gamedat", "txt", "txt-dir", "names"):
|
||
raise NameError(
|
||
f"{filename}: unrecognised WolfDawn document (kind={kind!r}). "
|
||
"Expected a strings-extract or names-extract JSON."
|
||
)
|
||
return parseDocument(data, filename)
|
||
|
||
|
||
def collectEntries(data):
|
||
"""Return the list of leaf {source, text} dicts for a WolfDawn document.
|
||
|
||
The returned dicts are live references into ``data`` so mutating ``text``
|
||
updates the document that gets written back out.
|
||
"""
|
||
kind = data.get("kind")
|
||
entries = []
|
||
|
||
if kind in ("map", "common"):
|
||
for scene in data.get("scenes") or []:
|
||
for line in scene.get("lines") or []:
|
||
entries.append(line)
|
||
elif kind == "db":
|
||
for group in data.get("groups") or []:
|
||
for line in group.get("lines") or []:
|
||
entries.append(line)
|
||
elif kind in ("gamedat", "txt"):
|
||
for line in data.get("lines") or []:
|
||
entries.append(line)
|
||
elif kind == "txt-dir":
|
||
for fileDoc in data.get("files") or []:
|
||
for line in fileDoc.get("lines") or []:
|
||
entries.append(line)
|
||
elif kind == "names":
|
||
for name in data.get("names") or []:
|
||
entries.append(name)
|
||
|
||
return entries
|
||
|
||
|
||
def parseDocument(data, filename):
|
||
"""Translate every translatable leaf entry and return [data, tokens, error]."""
|
||
global PBAR
|
||
totalTokens = [0, 0]
|
||
|
||
entries = collectEntries(data)
|
||
# Only translate entries that actually contain target-language text; the rest
|
||
# keep text == source so WolfDawn treats them as untouched on inject.
|
||
translatable = [
|
||
e for e in entries
|
||
if isinstance(e.get("source"), str) and re.search(LANGREGEX, e["source"])
|
||
]
|
||
|
||
with tqdm(bar_format=BAR_FORMAT, position=POSITION, total=len(translatable), leave=LEAVE) as pbar:
|
||
pbar.desc = filename
|
||
PBAR = pbar
|
||
if not translatable:
|
||
return [data, totalTokens, None]
|
||
|
||
sources = [e["source"] for e in translatable]
|
||
try:
|
||
response = translateAI(sources, [])
|
||
except Exception as e:
|
||
return [data, totalTokens, e]
|
||
|
||
translated, tokens = response[0], response[1]
|
||
totalTokens[0] += tokens[0]
|
||
totalTokens[1] += tokens[1]
|
||
|
||
# Write translations back (skip in estimate mode: translated == sources).
|
||
if not ESTIMATE and isinstance(translated, list) and len(translated) == len(translatable):
|
||
for entry, text in zip(translatable, translated):
|
||
if isinstance(text, str):
|
||
entry["text"] = text
|
||
|
||
return [data, totalTokens, None]
|
||
|
||
|
||
def getResultString(translatedData, translationTime, filename):
|
||
"""Format the per-file / total cost + status line for the console log."""
|
||
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) + "]"
|
||
)
|
||
timeString = Fore.BLUE + "[" + str(round(translationTime, 1)) + "s]"
|
||
|
||
if translatedData[2] is None:
|
||
return 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 (
|
||
filename + ": " + totalTokenstring + timeString
|
||
+ Fore.RED + " \u2717 " + errorString + Fore.RESET
|
||
)
|
||
|
||
|
||
def translateAI(text, history, history_ctx=None):
|
||
"""Thin wrapper around the shared translation entry point."""
|
||
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,
|
||
)
|