214 lines
7.1 KiB
Python
214 lines
7.1 KiB
Python
"""Language routing and output-validation helpers for prompt rewriting."""
|
||
|
||
import re
|
||
from typing import Callable, Dict, Tuple
|
||
|
||
|
||
_QUOTED_TEXT_PATTERN = re.compile(
|
||
r'"[^"\n]*"|“[^”\n]*”|‘[^’\n]*’|「[^」\n]*」|『[^』\n]*』'
|
||
)
|
||
def contains_han(text: str) -> bool:
|
||
return any("\u4e00" <= char <= "\u9fff" for char in text)
|
||
|
||
|
||
def contains_japanese(text: str) -> bool:
|
||
return any(
|
||
("\u3040" <= char <= "\u309f")
|
||
or ("\u30a0" <= char <= "\u30ff")
|
||
or ("\uff66" <= char <= "\uff9f")
|
||
for char in text
|
||
)
|
||
|
||
|
||
def detect_input_language(text: str) -> str:
|
||
"""Return ``ja``, ``zh``, or ``en`` using lightweight script detection."""
|
||
if contains_japanese(text):
|
||
return "ja"
|
||
if contains_han(text):
|
||
return "zh"
|
||
return "en"
|
||
|
||
|
||
def resolve_output_language(target_language: str, detected_language: str) -> str:
|
||
"""Resolve the supported output language while preserving current auto behavior."""
|
||
if target_language in {"en", "zh"}:
|
||
return target_language
|
||
if target_language != "auto":
|
||
raise ValueError(f"Unsupported target language: {target_language}")
|
||
|
||
# Output prompts currently support English and Simplified Chinese. Japanese
|
||
# input historically followed the Chinese prompt path, which is also the
|
||
# recommended output language for Wan models.
|
||
return "zh" if detected_language in {"ja", "zh"} else "en"
|
||
|
||
|
||
def append_prompt_suffix(prompt: str, suffix: str, language: str) -> str:
|
||
"""Append a quality suffix with punctuation appropriate for the language."""
|
||
if language not in {"en", "zh"}:
|
||
raise ValueError(f"Unsupported output language: {language}")
|
||
|
||
prompt = prompt.strip()
|
||
suffix = suffix.strip()
|
||
if not prompt:
|
||
return suffix
|
||
if not suffix:
|
||
return prompt
|
||
|
||
if prompt.endswith((",", ".", "!", "?", ";", ":")):
|
||
separator = " "
|
||
elif prompt.endswith(("，", "。", "！", "？", "；", "：")):
|
||
separator = ""
|
||
else:
|
||
separator = "，" if language == "zh" else ", "
|
||
|
||
return f"{prompt}{separator}{suffix}"
|
||
|
||
|
||
def output_requirements(language: str) -> str:
|
||
if language == "zh":
|
||
return (
|
||
"无论输入使用何种语言，都必须仅使用简体中文输出增强后的提示词。\n"
|
||
"只输出最终结果，不要解释。"
|
||
)
|
||
if language == "en":
|
||
return (
|
||
"Write the enhanced prompt entirely in English, regardless of the input language.\n"
|
||
"Output only the final result without explanation."
|
||
)
|
||
raise ValueError(f"Unsupported output language: {language}")
|
||
|
||
|
||
def build_enhancement_user_prompt(prompt: str, language: str) -> str:
|
||
return (
|
||
f"User Input: {prompt}\n\n"
|
||
f"Output Requirements:\n{output_requirements(language)}\n\n"
|
||
"Rewritten Prompt:"
|
||
)
|
||
|
||
|
||
def _source_quoted_literals(source_text: str) -> list[str]:
|
||
return sorted(
|
||
{match.group(0) for match in _QUOTED_TEXT_PATTERN.finditer(source_text)},
|
||
key=len,
|
||
reverse=True,
|
||
)
|
||
|
||
|
||
def _without_allowed_literals(text: str, source_text: str) -> str:
|
||
checked = text
|
||
for literal in _source_quoted_literals(source_text):
|
||
checked = checked.replace(literal, "")
|
||
return checked
|
||
|
||
|
||
def is_output_language_valid(
|
||
text: str,
|
||
language: str,
|
||
source_text: str = "",
|
||
) -> bool:
|
||
"""Check output script while allowing exact quoted literals from the input."""
|
||
checked = _without_allowed_literals(text.strip(), source_text)
|
||
if not checked:
|
||
return False
|
||
|
||
if language == "zh":
|
||
return contains_han(checked) and not contains_japanese(checked)
|
||
if language == "en":
|
||
return (
|
||
bool(re.search(r"[A-Za-z]", checked))
|
||
and not contains_han(checked)
|
||
and not contains_japanese(checked)
|
||
)
|
||
raise ValueError(f"Unsupported output language: {language}")
|
||
|
||
|
||
def protect_source_quoted_literals(
|
||
text: str,
|
||
source_text: str,
|
||
) -> Tuple[str, Dict[str, str]]:
|
||
"""Protect only quoted strings that appeared verbatim in the user's input."""
|
||
protected = text
|
||
placeholders: Dict[str, str] = {}
|
||
for literal in _source_quoted_literals(source_text):
|
||
if literal not in protected:
|
||
continue
|
||
placeholder = f"__PRESERVED_LITERAL_{len(placeholders)}__"
|
||
protected = protected.replace(literal, placeholder)
|
||
placeholders[placeholder] = literal
|
||
return protected, placeholders
|
||
|
||
|
||
def restore_protected_literals(
|
||
text: str,
|
||
placeholders: Dict[str, str],
|
||
) -> str:
|
||
restored = text
|
||
for placeholder, literal in placeholders.items():
|
||
restored = restored.replace(placeholder, literal)
|
||
return restored
|
||
|
||
|
||
def build_language_correction_messages(
|
||
enhanced_prompt: str,
|
||
language: str,
|
||
) -> list[dict[str, str]]:
|
||
if language == "zh":
|
||
system_prompt = (
|
||
"你是提示词语言规范化工具。将给定的增强提示词完整转换为简体中文。"
|
||
"不得添加、删除、概括或重新解释内容；必须保留所有主体、动作、时间变化、"
|
||
"镜头运动、构图、颜色、数量和约束。"
|
||
"所有形如__PRESERVED_LITERAL_n__的占位符必须原样保留。"
|
||
"只输出最终提示词正文，不要解释。"
|
||
)
|
||
elif language == "en":
|
||
system_prompt = (
|
||
"You are a prompt language normalizer. Convert the given enhanced prompt "
|
||
"entirely to English. Do not add, remove, summarize, or reinterpret content; "
|
||
"preserve every subject, action, temporal change, camera movement, "
|
||
"composition detail, color, quantity, and constraint. Preserve every "
|
||
"__PRESERVED_LITERAL_n__ placeholder exactly. Output only the final prompt "
|
||
"without explanation."
|
||
)
|
||
else:
|
||
raise ValueError(f"Unsupported output language: {language}")
|
||
|
||
return [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": f"Enhanced Prompt:\n{enhanced_prompt}"},
|
||
]
|
||
|
||
|
||
def ensure_output_language(
|
||
enhanced_prompt: str,
|
||
language: str,
|
||
source_text: str,
|
||
correction_callback: Callable[[list[dict[str, str]]], str],
|
||
) -> Tuple[str, bool]:
|
||
"""Return valid output, correcting it once only when the script mismatches."""
|
||
if is_output_language_valid(enhanced_prompt, language, source_text):
|
||
return enhanced_prompt, False
|
||
|
||
protected, placeholders = protect_source_quoted_literals(
|
||
enhanced_prompt,
|
||
source_text,
|
||
)
|
||
corrected = correction_callback(
|
||
build_language_correction_messages(protected, language)
|
||
).strip()
|
||
placeholders_preserved = all(
|
||
placeholder in corrected for placeholder in placeholders
|
||
)
|
||
corrected = restore_protected_literals(corrected, placeholders)
|
||
|
||
if not placeholders_preserved or not is_output_language_valid(
|
||
corrected,
|
||
language,
|
||
source_text,
|
||
):
|
||
raise RuntimeError(
|
||
f"Model failed to produce a valid '{language}' output "
|
||
"after one language-correction pass"
|
||
)
|
||
|
||
return corrected, True
|