f9fd86464d
Some checks failed
Windows / verify (push) Has been cancelled
`models/` 처럼 슬래시 없이 쓴 규칙은 깊이에 상관없이 같은 이름의 폴더를 전부 무시한다. 의도는 루트의 모델 캐시였는데 소스 패키지까지 같이 걸렸고, 그래서 asr/translator/manager/tiers/glossary/packs/speech_level 8개 파일이 한 번도 커밋된 적이 없다. 증상: Gitea 에서 clone 하면 `ModuleNotFoundError: No module named 'livesub.models'` 로 앱이 아예 임포트되지 않는다. 로컬 작업 트리에는 파일이 있으니 개발 중에는 절대 안 보인다. 방금 붙인 Windows CI VM 이 첫 검증에서 바로 잡아냈다. /models/, /adapters/, /data/ 로 루트에 고정하고 빠진 파일을 추가한다.
377 lines
15 KiB
Python
377 lines
15 KiB
Python
"""번역(MT) 래퍼.
|
||
|
||
두 백엔드를 지원한다.
|
||
|
||
* CTranslate2 — NLLB 계열 seq2seq. 가볍고 빠르며 추가학습이 쉽다.
|
||
* Transformers — Seed-X / Qwen3 같은 LLM. 품질이 좋고 LoRA 어댑터를 얹을 수 있다.
|
||
|
||
LLM 백엔드는 다시 두 가지 프롬프트 형식으로 갈린다. 같은 "LLM"이라도 요구사항이
|
||
정반대라 하나로 합치면 한쪽이 반드시 망가진다.
|
||
|
||
* PromptStyle.SEEDX — Seed-X 는 chat template 이 없는 **번역 전용 completion 모델**이다.
|
||
모델 카드가 정한 고정 문장과 끝의 `<언어코드>` 태그를 그대로 지켜야 하고,
|
||
그 밖의 지시문을 끼워 넣으면 학습 분포를 벗어난다. 따라서 용어집은 프롬프트가
|
||
아니라 **플레이스홀더 치환**으로 넣는다.
|
||
* PromptStyle.INSTRUCT — Qwen3 등. chat template 을 쓰고 용어집을 지시문으로 넣는다.
|
||
|
||
세 경우 모두 `Translator` 인터페이스를 따르므로 상위 레이어는 구분하지 않는다.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import abc
|
||
import logging
|
||
import re
|
||
import threading
|
||
|
||
from ..constants import LANGUAGES, models_dir
|
||
from .glossary import Glossary
|
||
from .speech_level import Politeness, SpeechLevel, fix_particles, prompt_instruction
|
||
from .tiers import MTBackend, MTSpec, PromptStyle
|
||
|
||
log = logging.getLogger(__name__)
|
||
|
||
#: 문장 끝에 LLM이 붙이곤 하는 군더더기를 걷어내기 위한 패턴
|
||
_TRAILING_NOISE = re.compile(r"^\s*(?:번역[:：]|Translation[:：])\s*", re.IGNORECASE)
|
||
|
||
|
||
def _restore_or_retranslate(out: str, replacements: list[str], retry) -> str:
|
||
"""자리표시자를 역어로 되돌린다. 하나라도 사라졌으면 용어집 없이 다시 번역.
|
||
|
||
번역 모델이 문장 일부를 통째로 누락하면 그 안의 자리표시자도 같이
|
||
사라진다. 그대로 복원하면 용어가 자막에서 증발하는데, 그건 용어집을
|
||
아예 안 쓴 것보다 나쁘다. 그래서 그럴 때는 보호 없이 한 번 더 번역해
|
||
원문 내용이 살아 있는 결과를 쓴다.
|
||
"""
|
||
missing = Glossary.missing_placeholders(out, replacements)
|
||
if not missing:
|
||
return Glossary.restore(out, replacements)
|
||
log.debug("자리표시자 %s 소실 — 용어집 없이 재번역", missing)
|
||
return retry()
|
||
|
||
|
||
def build_seedx_prompt(text: str, source_lang: str, target_lang: str) -> str:
|
||
"""Seed-X 모델 카드가 요구하는 정확한 프롬프트.
|
||
|
||
Translate the following English sentence into Chinese:
|
||
May the force be with you <zh>
|
||
|
||
끝의 `<언어코드>` 태그는 PPO 학습에 쓰인 것이라 **필수**다. 빠지면 번역
|
||
품질이 크게 흔들린다. 형식을 바꾸면 안 되므로 지시문이나 용어집 설명을
|
||
여기 끼워 넣지 말 것.
|
||
|
||
출처: https://huggingface.co/ByteDance-Seed/Seed-X-PPO-7B
|
||
"""
|
||
src = LANGUAGES[source_lang]["english"]
|
||
tgt = LANGUAGES[target_lang]["english"]
|
||
tag = LANGUAGES[target_lang]["seedx"]
|
||
return f"Translate the following {src} sentence into {tgt}:\n{text} <{tag}>"
|
||
|
||
|
||
class Translator(abc.ABC):
|
||
def __init__(self, spec: MTSpec, device: str = "cuda", device_index: int = 0) -> None:
|
||
self.spec = spec
|
||
self.device = device
|
||
self.device_index = device_index
|
||
self._lock = threading.Lock()
|
||
|
||
@property
|
||
@abc.abstractmethod
|
||
def loaded(self) -> bool: ...
|
||
|
||
@abc.abstractmethod
|
||
def load(self) -> None: ...
|
||
|
||
@abc.abstractmethod
|
||
def unload(self) -> None: ...
|
||
|
||
@abc.abstractmethod
|
||
def _translate(
|
||
self,
|
||
text: str,
|
||
source_lang: str,
|
||
target_lang: str,
|
||
glossary: Glossary | None,
|
||
speech_level: SpeechLevel,
|
||
source_politeness: Politeness,
|
||
) -> str: ...
|
||
|
||
def translate(
|
||
self,
|
||
text: str,
|
||
source_lang: str,
|
||
target_lang: str,
|
||
glossary: Glossary | None = None,
|
||
speech_level: SpeechLevel = SpeechLevel.POLITE,
|
||
source_politeness: Politeness = Politeness.UNKNOWN,
|
||
) -> str:
|
||
"""말투(speech_level)는 지시문을 이해하는 모델에만 프롬프트로 전달된다.
|
||
|
||
나머지 백엔드는 상위에서 `apply_speech_level` 후처리로 맞춘다.
|
||
"""
|
||
text = text.strip()
|
||
if not text:
|
||
return ""
|
||
if source_lang == target_lang:
|
||
return text
|
||
out = self._translate(
|
||
text, source_lang, target_lang, glossary, speech_level, source_politeness
|
||
)
|
||
out = _TRAILING_NOISE.sub("", out).strip()
|
||
if target_lang == "ko":
|
||
# 용어집이 단어를 바꿔치기하면 앞말 받침이 달라져 조사가 어긋난다.
|
||
# (자기장를 -> 자기장을)
|
||
out = fix_particles(out)
|
||
return out
|
||
|
||
|
||
class CTranslate2Translator(Translator):
|
||
"""NLLB 계열. 용어집은 플레이스홀더 보호 방식으로 적용한다."""
|
||
|
||
def __init__(self, spec: MTSpec, device: str = "cuda", device_index: int = 0) -> None:
|
||
super().__init__(spec, device, device_index)
|
||
self._model = None
|
||
self._tokenizer = None
|
||
|
||
@property
|
||
def loaded(self) -> bool:
|
||
return self._model is not None
|
||
|
||
def load(self) -> None:
|
||
if self._model is not None:
|
||
return
|
||
import ctranslate2
|
||
from huggingface_hub import snapshot_download
|
||
from transformers import AutoTokenizer
|
||
|
||
local = snapshot_download(self.spec.repo, cache_dir=str(models_dir()))
|
||
compute = self.spec.compute_type
|
||
if self.device == "cpu" and "float16" in compute:
|
||
compute = "int8"
|
||
log.info("번역 모델 로드: %s (%s/%s)", self.spec.repo, self.device, compute)
|
||
self._model = ctranslate2.Translator(
|
||
local,
|
||
device=self.device,
|
||
device_index=self.device_index,
|
||
compute_type=compute,
|
||
)
|
||
self._tokenizer = AutoTokenizer.from_pretrained(local, cache_dir=str(models_dir()))
|
||
|
||
def unload(self) -> None:
|
||
with self._lock:
|
||
self._model = None
|
||
self._tokenizer = None
|
||
|
||
def _run_model(self, text: str, source_lang: str, target_lang: str) -> str:
|
||
"""용어집 처리를 뺀 순수 번역 한 번."""
|
||
src_code = LANGUAGES[source_lang]["nllb"]
|
||
tgt_code = LANGUAGES[target_lang]["nllb"]
|
||
with self._lock:
|
||
self._tokenizer.src_lang = src_code
|
||
tokens = self._tokenizer.convert_ids_to_tokens(
|
||
self._tokenizer.encode(text, truncation=True, max_length=512)
|
||
)
|
||
results = self._model.translate_batch(
|
||
[tokens],
|
||
target_prefix=[[tgt_code]],
|
||
beam_size=2,
|
||
max_decoding_length=512,
|
||
repetition_penalty=1.1,
|
||
)
|
||
hyp = results[0].hypotheses[0]
|
||
if hyp and hyp[0] == tgt_code:
|
||
hyp = hyp[1:]
|
||
return self._tokenizer.decode(
|
||
self._tokenizer.convert_tokens_to_ids(hyp), skip_special_tokens=True
|
||
)
|
||
|
||
def _translate(self, text, source_lang, target_lang, glossary,
|
||
speech_level=SpeechLevel.POLITE, source_politeness=Politeness.UNKNOWN):
|
||
# seq2seq 는 프롬프트가 없어 말투를 지시할 수 없다. 후처리로 맞춘다.
|
||
if self._model is None:
|
||
self.load()
|
||
assert self._model is not None and self._tokenizer is not None
|
||
|
||
if glossary is None or not len(glossary):
|
||
return self._run_model(text, source_lang, target_lang)
|
||
|
||
protected, replacements = glossary.protect(text, target_lang)
|
||
if not replacements:
|
||
return self._run_model(text, source_lang, target_lang)
|
||
|
||
out = self._run_model(protected, source_lang, target_lang)
|
||
return _restore_or_retranslate(
|
||
out, replacements,
|
||
lambda: self._run_model(text, source_lang, target_lang),
|
||
)
|
||
|
||
|
||
class LLMTranslator(Translator):
|
||
"""Seed-X / Qwen3 등 LLM 백엔드. 용어집은 프롬프트로 직접 지시한다."""
|
||
|
||
def __init__(self, spec: MTSpec, device: str = "cuda", device_index: int = 0,
|
||
lora_path: str = "") -> None:
|
||
super().__init__(spec, device, device_index)
|
||
self.lora_path = lora_path
|
||
self._model = None
|
||
self._tokenizer = None
|
||
|
||
@property
|
||
def loaded(self) -> bool:
|
||
return self._model is not None
|
||
|
||
def load(self) -> None:
|
||
if self._model is not None:
|
||
return
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
||
dtype = torch.float16 if self.device == "cuda" else torch.float32
|
||
log.info("번역 LLM 로드: %s (%s)", self.spec.repo, self.device)
|
||
self._tokenizer = AutoTokenizer.from_pretrained(
|
||
self.spec.repo, cache_dir=str(models_dir())
|
||
)
|
||
try:
|
||
self._model = AutoModelForCausalLM.from_pretrained(
|
||
self.spec.repo,
|
||
cache_dir=str(models_dir()),
|
||
torch_dtype=dtype,
|
||
device_map={"": self.device_index} if self.device == "cuda" else "cpu",
|
||
)
|
||
except (ImportError, RuntimeError, ValueError) as exc:
|
||
if self.spec.compute_type == "int4":
|
||
# AWQ 가중치는 전용 커널이 있어야 읽힌다. 메시지가 모호해서
|
||
# 사용자가 원인을 못 찾는 일이 잦으므로 해결 방법을 같이 알려준다.
|
||
raise RuntimeError(
|
||
f"양자화 번역 모델({self.spec.repo})을 불러오지 못했습니다.\n"
|
||
"AWQ 커널이 필요합니다: pip install autoawq\n"
|
||
"설치가 어려우면 '모델' 화면에서 다른 티어를 선택하세요.\n"
|
||
f"원인: {exc}"
|
||
) from exc
|
||
raise
|
||
if self.lora_path:
|
||
from peft import PeftModel
|
||
|
||
log.info("LoRA 어댑터 적용: %s", self.lora_path)
|
||
self._model = PeftModel.from_pretrained(self._model, self.lora_path)
|
||
self._model.eval()
|
||
|
||
def unload(self) -> None:
|
||
with self._lock:
|
||
self._model = None
|
||
self._tokenizer = None
|
||
try:
|
||
import torch
|
||
|
||
torch.cuda.empty_cache()
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
|
||
def _build_prompt(self, text, source_lang, target_lang, glossary,
|
||
speech_level=SpeechLevel.POLITE,
|
||
source_politeness=Politeness.UNKNOWN) -> str:
|
||
if self.spec.prompt_style is PromptStyle.SEEDX:
|
||
# Seed-X 는 지시문을 못 알아듣는다. 말투는 후처리로만 맞춘다.
|
||
return build_seedx_prompt(text, source_lang, target_lang)
|
||
return self._build_instruct_prompt(
|
||
text, source_lang, target_lang, glossary, speech_level, source_politeness
|
||
)
|
||
|
||
def _build_instruct_prompt(self, text, source_lang, target_lang, glossary,
|
||
speech_level=SpeechLevel.POLITE,
|
||
source_politeness=Politeness.UNKNOWN) -> str:
|
||
"""지시문을 이해하는 모델용. 가능하면 chat template 을 태운다."""
|
||
src = LANGUAGES[source_lang]["english"]
|
||
tgt = LANGUAGES[target_lang]["english"]
|
||
hint = glossary.prompt_hint(text, target_lang) if glossary else ""
|
||
parts = [
|
||
f"Translate the following {src} text into {tgt}. "
|
||
"It is a live spoken line from a game or broadcast, so keep the tone "
|
||
"casual and natural. Output only the translation."
|
||
]
|
||
if target_lang == "ko":
|
||
parts.append(prompt_instruction(speech_level, source_politeness))
|
||
if hint:
|
||
parts.append(hint)
|
||
parts.append(f"{src}: {text}\n{tgt}:")
|
||
instruction = "\n\n".join(parts)
|
||
|
||
template = getattr(self._tokenizer, "chat_template", None)
|
||
if not template:
|
||
return instruction
|
||
# Qwen3 는 기본적으로 thinking 모드가 켜져 실시간 자막에는 너무 느리다.
|
||
try:
|
||
return self._tokenizer.apply_chat_template(
|
||
[{"role": "user", "content": instruction}],
|
||
tokenize=False,
|
||
add_generation_prompt=True,
|
||
enable_thinking=False,
|
||
)
|
||
except TypeError:
|
||
# enable_thinking 을 모르는 템플릿
|
||
return self._tokenizer.apply_chat_template(
|
||
[{"role": "user", "content": instruction}],
|
||
tokenize=False,
|
||
add_generation_prompt=True,
|
||
)
|
||
|
||
def _generate(self, prompt: str) -> str:
|
||
"""프롬프트 하나를 돌려 첫 줄만 돌려준다."""
|
||
import torch
|
||
|
||
with self._lock:
|
||
inputs = self._tokenizer(prompt, return_tensors="pt").to(self._model.device)
|
||
with torch.inference_mode():
|
||
generated = self._model.generate(
|
||
**inputs,
|
||
max_new_tokens=256,
|
||
do_sample=False,
|
||
num_beams=1,
|
||
repetition_penalty=1.05,
|
||
pad_token_id=self._tokenizer.eos_token_id,
|
||
)
|
||
new_tokens = generated[0][inputs["input_ids"].shape[-1] :]
|
||
out = self._tokenizer.decode(new_tokens, skip_special_tokens=True)
|
||
return out.strip().split("\n")[0]
|
||
|
||
def _translate(self, text, source_lang, target_lang, glossary,
|
||
speech_level=SpeechLevel.POLITE, source_politeness=Politeness.UNKNOWN):
|
||
if self._model is None:
|
||
self.load()
|
||
assert self._model is not None and self._tokenizer is not None
|
||
|
||
def run(source: str, gloss) -> str:
|
||
return self._generate(
|
||
self._build_prompt(
|
||
source, source_lang, target_lang, gloss,
|
||
speech_level, source_politeness,
|
||
)
|
||
)
|
||
|
||
# Seed-X 는 지시문을 못 알아들으므로 용어집을 자리표시자로 보호한다.
|
||
# 지시문을 이해하는 모델은 프롬프트로 넣으므로 보호가 필요 없다.
|
||
needs_placeholder = (
|
||
self.spec.prompt_style is PromptStyle.SEEDX
|
||
and glossary is not None
|
||
and len(glossary)
|
||
)
|
||
if not needs_placeholder:
|
||
return run(text, glossary)
|
||
|
||
protected, replacements = glossary.protect(text, target_lang)
|
||
if not replacements:
|
||
return run(text, glossary)
|
||
|
||
out = run(protected, None)
|
||
return _restore_or_retranslate(
|
||
out, replacements, lambda: run(text, None)
|
||
)
|
||
|
||
|
||
def create_translator(
|
||
spec: MTSpec, device: str = "cuda", device_index: int = 0, lora_path: str = ""
|
||
) -> Translator:
|
||
if spec.backend is MTBackend.CTRANSLATE2:
|
||
return CTranslate2Translator(spec, device, device_index)
|
||
return LLMTranslator(spec, device, device_index, lora_path)
|