89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import urllib.error
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
from wov_sdk.models import InvokeRequest, InvokeResponse, NodeManifest
|
|
from wov_sdk.server import run_node
|
|
|
|
CHUNK_SIZE = 20
|
|
|
|
|
|
def translate_lines(lines: list[str], params: dict) -> list[str]:
|
|
api_base = os.getenv(
|
|
"LLM_API_BASE",
|
|
"http://192.168.123.70:8080/v1/chat/completions",
|
|
)
|
|
api_key = os.getenv("LLM_API_KEY", "")
|
|
model = str(params.get("model") or os.getenv("LLM_MODEL", "default"))
|
|
target_language = str(params.get("target_language", "zh-CN"))
|
|
system_prompt = (
|
|
"你是专业字幕翻译。将用户提供的日文字幕翻译为"
|
|
f"{target_language}。只返回译文，保持行数和顺序，不要添加解释。"
|
|
)
|
|
translated: list[str] = []
|
|
for start in range(0, len(lines), CHUNK_SIZE):
|
|
chunk = lines[start : start + CHUNK_SIZE]
|
|
body = {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": "\n".join(chunk)},
|
|
],
|
|
}
|
|
headers = {"Content-Type": "application/json"}
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
request = urllib.request.Request(
|
|
api_base,
|
|
data=json.dumps(body).encode("utf-8"),
|
|
headers=headers,
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(request, timeout=120) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
content = payload["choices"][0]["message"]["content"]
|
|
translated.extend(
|
|
[line.strip() for line in content.splitlines() if line.strip()]
|
|
)
|
|
return translated
|
|
|
|
|
|
def invoke(request: InvokeRequest) -> InvokeResponse:
|
|
srt_uri = request.inputs.get("srt_uri")
|
|
if not srt_uri:
|
|
return InvokeResponse(status="failed", error="srt_uri is required")
|
|
|
|
srt_path = Path(srt_uri)
|
|
if not srt_path.is_file():
|
|
return InvokeResponse(status="failed", error="srt file not found")
|
|
|
|
lines = srt_path.read_text(encoding="utf-8").splitlines()
|
|
text_indices = list(range(2, len(lines), 4))
|
|
source_lines = [lines[index] for index in text_indices]
|
|
translated_lines = translate_lines(source_lines, request.params)
|
|
translated_lines = translated_lines[: len(source_lines)]
|
|
translated_lines += [""] * max(0, len(source_lines) - len(translated_lines))
|
|
for index, text_index in enumerate(text_indices):
|
|
lines[text_index] = translated_lines[index]
|
|
|
|
output_dir = Path(request.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
output_path = output_dir / "cn.srt"
|
|
output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
return InvokeResponse(status="completed", outputs={"cn_srt_uri": str(output_path)})
|
|
|
|
|
|
def main() -> None:
|
|
manifest_path = Path(__file__).resolve().parent.parent / "node.manifest.json"
|
|
with open(manifest_path, "r", encoding="utf-8") as f:
|
|
manifest = NodeManifest.from_dict(json.load(f))
|
|
run_node(manifest, invoke)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|