epub-translator 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

epub_translator/template.py

@@ -0,0 +1,50 @@
+import re
+
+from typing import Tuple, Callable
+from pathlib import Path
+from jinja2 import select_autoescape, Environment, BaseLoader, TemplateNotFound
+
+
+def create_env(dir_path: Path) -> Environment:
+  return Environment(
+    loader=_DSLoader(dir_path),
+    autoescape=select_autoescape(),
+    trim_blocks=True,
+    keep_trailing_newline=True,
+  )
+
+_LoaderResult = Tuple[str, str | None, Callable[[], bool] | None]
+
+class _DSLoader(BaseLoader):
+  def __init__(self, dir_path: Path):
+    super().__init__()
+    self._dir_path: Path = dir_path
+
+  def get_source(self, _: Environment, template: str) -> _LoaderResult:
+    template = self._norm_template(template)
+    target_path = (self._dir_path / template).resolve()
+
+    if not target_path.exists():
+      raise TemplateNotFound(f"cannot find {template}")
+
+    return self._get_source_with_path(target_path)
+
+  def _norm_template(self, template: str) -> str:
+    if bool(re.match(r"^\.+/", template)):
+      raise TemplateNotFound(f"invalid path {template}")
+
+    template = re.sub(r"^/", "", template)
+    template = re.sub(r"\.jinja$", "", template, flags=re.IGNORECASE)
+    template = f"{template}.jinja"
+
+    return template
+
+  def _get_source_with_path(self, path: Path) -> _LoaderResult:
+    mtime = path.stat().st_mtime
+    with open(path, "r", encoding="utf-8") as f:
+      source = f.read()
+
+    def is_updated() -> bool:
+      return mtime == path.stat().st_mtime
+
+    return source, path, is_updated

epub_translator/translation/__init__.py

@@ -0,0 +1,2 @@
+from .types import *
+from .translation import translate, ProgressReporter

epub_translator/translation/chunk.py

@@ -0,0 +1,120 @@
+from dataclasses import dataclass
+from typing import Iterator, Iterable, Generator
+from hashlib import sha512
+from ..llm import LLM
+from .types import Fragment
+
+
+@dataclass
+class Chunk:
+  index: int
+  hash: bytes
+  head: list[str]
+  body: list[str]
+  tail: list[str]
+  tokens_count: int
+
+@dataclass
+class ChunkRange:
+  index: int
+  head_remain_tokens: int
+  tail_remain_tokens: int
+  head_index: int
+  body_index: int
+  tail_index: int
+  fragments_count: int
+  tokens_count: int
+
+  def match(self, index: int) -> bool:
+    return self.head_index <= index < self.head_index + self.fragments_count
+
+def match_fragments(
+        llm: LLM,
+        chunk_ranges_iter: Iterator[ChunkRange],
+        fragments_iter: Iterator[Fragment],
+      ) -> Generator[Chunk, None, None]:
+
+  for range, texts in _match_range_and_texts(
+    chunk_range_iter=chunk_ranges_iter,
+    fragments_iter=fragments_iter,
+  ):
+    head_length = range.body_index - range.head_index
+    body_length = range.tail_index - range.body_index
+    head = texts[:head_length]
+    body = texts[head_length:head_length + body_length]
+    tail = texts[head_length + body_length:]
+
+    hash = _hash_texts_list((head, body, tail))
+    head = _crop_extra_texts(llm, head, True, range.head_remain_tokens)
+    tail = _crop_extra_texts(llm, tail, False, range.tail_remain_tokens)
+
+    yield Chunk(
+      hash=hash,
+      head=head,
+      body=body,
+      tail=tail,
+      index=range.index,
+      tokens_count=range.tokens_count,
+    )
+
+def _match_range_and_texts(
+      chunk_range_iter: Iterator[ChunkRange],
+      fragments_iter: Iterator[Fragment],
+    ) -> Generator[tuple[ChunkRange, list[str]], None, None]:
+
+  next_chunk_range: ChunkRange | None = None
+  matched_chunk_ranges: list[tuple[ChunkRange, list[str]]] = []
+
+  for index, fragment in enumerate(fragments_iter):
+    while True:
+      if next_chunk_range is None:
+        next_chunk_range = next(chunk_range_iter, None)
+        if next_chunk_range is None:
+          break
+      if not next_chunk_range.match(index):
+        break
+      matched_chunk_ranges.append((next_chunk_range, []))
+      next_chunk_range = None
+
+    if matched_chunk_ranges:
+      next_matched_chunks: list[tuple[ChunkRange, list[str]]] = []
+      for chunk_range, texts in matched_chunk_ranges:
+        if chunk_range.match(index):
+          texts.append(fragment.text)
+          next_matched_chunks.append((chunk_range, texts))
+        else:
+          yield chunk_range, texts
+      matched_chunk_ranges = next_matched_chunks
+
+  yield from matched_chunk_ranges
+
+def _hash_texts_list(texts_iterable: Iterable[list[str]]) -> bytes:
+  is_first = True
+  m = sha512()
+  for texts in texts_iterable:
+    for text in texts:
+      if is_first:
+        is_first = False
+      else:
+        m.update(b"\x00")
+      m.update(text.encode("utf-8"))
+  return m.digest()
+
+def _crop_extra_texts(llm: LLM, texts: list[str], crop_left: bool, remain_tokens_count: int):
+  tokens_list: list[list[int]] = [llm.encode_tokens(text) for text in texts]
+  remain_texts: list[str] = []
+
+  for tokens in (reversed(tokens_list) if crop_left else tokens_list):
+    tokens_count = len(tokens)
+    if remain_tokens_count >= tokens_count:
+      remain_tokens_count -= tokens_count
+      remain_texts.append(llm.decode_tokens(tokens))
+      if remain_tokens_count == 0:
+        break
+    else:
+      remain_tokens = tokens[-remain_tokens_count:] if crop_left else tokens[:remain_tokens_count]
+      remain_texts.append(llm.decode_tokens(remain_tokens))
+
+  if crop_left:
+    remain_texts.reverse()
+  return remain_texts

epub_translator/translation/splitter.py

@@ -0,0 +1,77 @@
+from typing import Iterator, Generator
+from resource_segmentation import split, Resource, Segment
+
+from ..llm import LLM
+from .types import Fragment
+from .chunk import ChunkRange
+
+
+def split_into_chunks(llm: LLM, fragments_iter: Iterator[Fragment], max_chunk_tokens_count: int):
+  for index, group in enumerate(split(
+    resources=_gen_resources(llm, fragments_iter),
+    max_segment_count=max_chunk_tokens_count,
+    gap_rate=0.15,
+    tail_rate=0.5,
+  )):
+    head_index: int
+    tail_index: int
+    fragments_count: int
+    body_index, body_end_index, body_tokens_count = _range_of_group_part(group.body)
+
+    if group.head:
+      head_index, head_end_index, _ = _range_of_group_part(group.head)
+      assert head_end_index + 1 == body_index, "Head must be continuous with body"
+    else:
+      head_index = body_index
+
+    if group.tail:
+      tail_index, tail_end_index, _ = _range_of_group_part(group.tail)
+      fragments_count = tail_end_index - head_index + 1
+      assert body_end_index + 1 == tail_index, "Body must be continuous with tail"
+    else:
+      tail_index = body_end_index + 1
+      fragments_count = tail_index - head_index
+
+    yield ChunkRange(
+      index=index,
+      head_remain_tokens=group.head_remain_count,
+      tail_remain_tokens=group.tail_remain_count,
+      head_index=head_index,
+      body_index=body_index,
+      tail_index=tail_index,
+      fragments_count=fragments_count,
+      tokens_count=body_tokens_count,
+    )
+
+def _gen_resources(llm: LLM, fragments_iter: Iterator[Fragment]) -> Generator[Resource[int], None, None]:
+  for index, fragment in enumerate(fragments_iter):
+    yield Resource(
+      count=llm.count_tokens_count(fragment.text),
+      start_incision=fragment.start_incision,
+      end_incision=fragment.end_incision,
+      payload=index,
+    )
+
+def _range_of_group_part(target: list[Resource[int] | Segment[int]]) -> tuple[int, int]:
+  start_index: int | None = None
+  previous_index: int = 0
+  tokens_count: int = 0
+  for resource in _iter_group_part(target):
+    index = resource.payload
+    if start_index is None:
+      start_index = index
+    else:
+      assert index == previous_index + 1, "Resources in group part must be continuous"
+    previous_index = index
+    tokens_count += resource.count
+
+  assert start_index is not None, "Group part must contain at least one resource"
+  return start_index, previous_index, tokens_count
+
+def _iter_group_part(target: list[Resource[int] | Segment[int]]) -> Generator[Resource[int], None, None]:
+  for item in target:
+    if isinstance(item, Resource):
+      yield item
+    elif isinstance(item, Segment):
+      for resource in item.resources:
+        yield resource

epub_translator/translation/store.py

@@ -0,0 +1,37 @@
+from shutil import rmtree
+from pathlib import Path
+from typing import Iterator
+from .utils import clean_spaces
+
+
+class Store:
+  def __init__(self, directory: Path):
+    self._directory = directory
+
+  def get(self, chunk_hash: bytes) -> list[str] | None:
+    file_path = self._file_path(chunk_hash)
+    if not file_path.exists() or not file_path.is_file():
+      return None
+    with file_path.open("r", encoding="utf-8") as file:
+      return list(line for line in file if line.strip())
+
+  def put(self, chunk_hash: bytes, lines_iter: Iterator[str]):
+    file_path = self._file_path(chunk_hash)
+    if file_path.exists():
+      if file_path.is_file():
+        file_path.unlink()
+      else:
+        rmtree(file_path)
+
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+    with file_path.open("w", encoding="utf-8") as file:
+      is_first_line = True
+      for line in lines_iter:
+        if is_first_line:
+          is_first_line = False
+        else:
+          file.write("\n")
+        file.write(clean_spaces(line))
+
+  def _file_path(self, chunk_hash: bytes) -> Path:
+    return self._directory / f"{chunk_hash.hex()}.chunk"

epub_translator/translation/translation.py

@@ -0,0 +1,192 @@
+from typing import Callable, Iterator, Generator
+from pathlib import Path
+from concurrent.futures import as_completed, ThreadPoolExecutor
+from xml.etree.ElementTree import Element
+
+from ..llm import LLM
+from ..xml import encode_friendly
+
+from .types import Fragment, Language
+from .store import Store
+from .splitter import split_into_chunks
+from .chunk import match_fragments, Chunk
+from .utils import is_empty, clean_spaces
+
+
+ProgressReporter = Callable[[float], None]
+
+def translate(
+      llm: LLM,
+      gen_fragments_iter: Callable[[], Iterator[Fragment]],
+      cache_path: Path | None,
+      target_language: Language,
+      user_prompt: str | None,
+      max_chunk_tokens_count: int,
+      max_threads_count: int,
+      report_progress: ProgressReporter,
+    ) -> Generator[str, None, None]:
+
+  if user_prompt is not None:
+    user_prompt = _normalize_user_input(user_prompt.splitlines())
+
+  store = Store(cache_path) if cache_path else None
+  chunk_ranges = list(split_into_chunks(
+    llm=llm,
+    fragments_iter=gen_fragments_iter(),
+    max_chunk_tokens_count=max_chunk_tokens_count,
+  ))
+  with ThreadPoolExecutor(max_workers=max_threads_count) as executor:
+    futures = [
+      executor.submit(lambda chunk=chunk: (chunk, _translate_chunk(
+        llm=llm,
+        store=store,
+        chunk=chunk,
+        target_language=target_language,
+        user_prompt=user_prompt,
+      )))
+      for chunk in match_fragments(
+        llm=llm,
+        chunk_ranges_iter=iter(chunk_ranges),
+        fragments_iter=gen_fragments_iter(),
+      )
+    ]
+    yield from _sort_translated_texts_by_chunk(
+      target=(f.result() for f in as_completed(futures)),
+      total_tokens_count=sum(chunk.tokens_count for chunk in chunk_ranges),
+      report_progress=report_progress,
+    )
+
+def _sort_translated_texts_by_chunk(
+      target: Iterator[tuple[Chunk, list[str]]],
+      total_tokens_count: int,
+      report_progress: ProgressReporter,
+    ) -> Iterator[list[str]]:
+
+  buffer: list[tuple[Chunk, list[str]]] = []
+  wanna_next_index: int = 0
+  translated_tokens_count: int = 0
+
+  for chunk, translated_texts in target:
+    buffer.append((chunk, translated_texts))
+    if wanna_next_index == chunk.index:
+      buffer.sort(key=lambda e: e[0].index)
+      to_clear: list[list[str]] = []
+
+      for chunk, translated_texts in buffer:
+        if chunk.index > wanna_next_index:
+          break
+        to_clear.append(translated_texts)
+        if chunk.index == wanna_next_index:
+          wanna_next_index += 1
+
+      if to_clear:
+        buffer = buffer[len(to_clear):]
+        for translated_texts in to_clear:
+          yield from translated_texts
+
+    translated_tokens_count += chunk.tokens_count
+    report_progress(float(translated_tokens_count) / total_tokens_count)
+
+def _translate_chunk(
+      llm: LLM,
+      store: Store,
+      chunk: Chunk,
+      target_language: Language,
+      user_prompt: str | None,
+    ) -> list[str]:
+
+    translated_texts: list[str] | None = None
+    if store is not None:
+      translated_texts = store.get(chunk.hash)
+
+    if translated_texts is None:
+      translated_texts = _translate_texts(
+        llm=llm,
+        texts=chunk.head + chunk.body + chunk.tail,
+        target_language=target_language,
+        user_prompt=user_prompt,
+      )
+    if store is not None:
+      store.put(chunk.hash, translated_texts)
+
+    head_length = len(chunk.head)
+    translated_texts = translated_texts[head_length:head_length + len(chunk.body)]
+
+    return translated_texts
+
+def _translate_texts(
+      llm: LLM,
+      texts: list[str],
+      target_language: Language,
+      user_prompt: str | None,
+    ) -> list[str]:
+
+  original_text = _normalize_user_input(texts)
+  if original_text is None:
+    return [""] * len(texts)
+
+  user_data = original_text
+  if user_prompt is not None:
+    user_data = f"<rules>{user_prompt}</rules>\n\n{original_text}"
+
+  translated_text = llm.request_text(
+    template_name="translate",
+    text_tag="TXT",
+    user_data=user_data,
+    parser=lambda r: r,
+    params={
+      "target_language": target_language.value,
+      "user_prompt": user_prompt,
+    },
+  )
+  request_element = Element("request")
+
+  for i, fragment in enumerate(texts):
+    fragment_element = Element("fragment", attrib={
+      "id": str(i + 1),
+    })
+    fragment_element.text = clean_spaces(fragment)
+    request_element.append(fragment_element)
+
+  request_element_text = encode_friendly(request_element)
+  request_text = f"```XML\n{request_element_text}\n```\n\n{translated_text}"
+
+  return llm.request_xml(
+    template_name="format",
+    user_data=request_text,
+    params={ "target_language": target_language.value },
+    parser=lambda r: _parse_translated_response(r, len(texts)),
+  )
+
+def _parse_translated_response(resp_element: Element, sources_count: int) -> list[str]:
+  translated_fragments = [""] * sources_count
+  for fragment_element in resp_element:
+    if fragment_element.text is None:
+      continue
+    id = fragment_element.get("id", None)
+    if id is None:
+      continue
+    index = int(id) - 1
+    if index < 0 or index >= len(translated_fragments):
+      raise ValueError(f"invalid fragment id: {id}")
+    translated_fragments[index] = fragment_element.text.strip()
+
+  return translated_fragments
+
+def _normalize_user_input(user_lines: list[str]) -> str | None:
+  empty_lines_count: int = 0
+  lines: list[str] = []
+  for line in user_lines:
+    if is_empty(line):
+      empty_lines_count += 1
+    else:
+      if lines:
+        if empty_lines_count >= 2:
+          lines.append("")
+          lines.append("")
+        elif empty_lines_count == 1:
+          lines.append("")
+      lines.append(clean_spaces(line))
+  if not lines:
+    return None
+  return "\n".join(lines)

epub_translator/translation/types.py

@@ -0,0 +1,23 @@
+from enum import Enum
+from dataclasses import dataclass
+from resource_segmentation import Incision
+
+
+@dataclass
+class Fragment:
+  text: str
+  start_incision: Incision
+  end_incision: Incision
+
+class Language(Enum):
+  SIMPLIFIED_CHINESE = "简体中文"
+  TRADITIONAL_CHINESE = "繁体中文"
+  ENGLISH = "英语"
+  FRENCH = "法语"
+  GERMAN = "德语"
+  SPANISH = "西班牙语"
+  RUSSIAN = "俄语"
+  ITALIAN = "意大利语"
+  PORTUGUESE = "葡萄牙语"
+  JAPANESE = "日语"
+  KOREAN = "韩语"

epub_translator/translation/utils.py

@@ -0,0 +1,11 @@
+import re
+
+
+_EMPTY_LINE = re.compile(r"^\s*$")
+_SPACE = re.compile(r"\s+")
+
+def is_empty(text: str) -> bool:
+  return bool(_EMPTY_LINE.match(text))
+
+def clean_spaces(text: str) -> str:
+  return _SPACE.sub(" ", text.strip())

epub_translator/translator.py

@@ -0,0 +1,169 @@
+from os import PathLike
+from pathlib import Path
+from tempfile import mkdtemp
+from shutil import rmtree
+
+from .llm import LLM
+from .epub import HTMLFile
+from .zip_context import ZipContext
+from .translation import translate as _translate, Fragment, Incision, Language, ProgressReporter
+
+
+def translate(
+      llm: LLM,
+      source_path: PathLike,
+      translated_path: PathLike,
+      target_language: Language,
+      user_prompt: str | None = None,
+      working_path: PathLike | None = None,
+      max_chunk_tokens_count: int = 3000,
+      max_threads_count: int = 1,
+      report_progress: ProgressReporter | None = None,
+    ) -> None:
+
+  source_path = Path(source_path)
+  translated_path = Path(translated_path)
+  working_path = Path(working_path) if working_path else None
+  report_progress = report_progress or (lambda _: None)
+
+  _Translator(
+    llm=llm,
+    target_language=target_language,
+    user_prompt=user_prompt,
+    max_chunk_tokens_count=max_chunk_tokens_count,
+    max_threads_count=max_threads_count,
+    report_progress=report_progress,
+  ).do(
+    source_path=source_path,
+    translated_path=translated_path,
+    working_path=working_path,
+  )
+
+class _Translator:
+  def __init__(
+        self,
+        llm: LLM,
+        target_language: Language,
+        user_prompt: str | None,
+        max_chunk_tokens_count: int,
+        max_threads_count: int,
+        report_progress: ProgressReporter,
+      ) -> None:
+
+    self._llm: LLM = llm
+    self._target_language: Language = target_language
+    self._user_prompt: str | None = user_prompt
+    self._max_chunk_tokens_count: int = max_chunk_tokens_count
+    self._max_threads_count: int = max_threads_count
+    self._report_progress: ProgressReporter = report_progress
+
+  def do(self, source_path: Path, translated_path: Path, working_path: Path | None) -> None:
+    is_temp_workspace = not bool(working_path)
+    working_path = working_path or Path(mkdtemp())
+    try:
+      temp_dir = _clean_path(working_path / "temp")
+      temp_dir.mkdir(parents=True, exist_ok=True)
+      cache_path = working_path / "cache"
+
+      context = ZipContext(
+        epub_path=Path(source_path),
+        temp_dir=temp_dir,
+      )
+      context.replace_ncx(lambda texts: self._translate_ncx(
+        texts=texts,
+        cache_path=cache_path,
+        report_progress=lambda p: self._report_progress(p * 0.1)),
+      )
+      self._translate_spine(
+        context=context,
+        cache_path=cache_path,
+        report_progress=lambda p: self._report_progress(0.1 + p * 0.8),
+      )
+      context.archive(translated_path)
+      self._report_progress(1.0)
+
+    finally:
+      if is_temp_workspace:
+        rmtree(working_path, ignore_errors=True)
+
+  def _translate_ncx(self, texts: list[str], cache_path: Path, report_progress: ProgressReporter) -> list[str]:
+    return list(_translate(
+      llm=self._llm,
+      cache_path=cache_path,
+      max_chunk_tokens_count=self._max_chunk_tokens_count,
+      max_threads_count=1,
+      target_language=self._target_language,
+      user_prompt=self._user_prompt,
+      report_progress=report_progress,
+      gen_fragments_iter=lambda: (
+        Fragment(
+          text=text,
+          start_incision=Incision.IMPOSSIBLE,
+          end_incision=Incision.IMPOSSIBLE,
+        )
+        for text in texts
+      ),
+    ))
+
+  def _translate_spine(self, context: ZipContext, cache_path: Path, report_progress: ProgressReporter):
+    spine_paths_iter = iter(list(context.search_spine_paths()))
+    spine_file: HTMLFile | None = None
+    translated_texts: list[str] = []
+    translated_count: int = 0
+
+    for translated_text in _translate(
+      llm=self._llm,
+      gen_fragments_iter=lambda: _gen_fragments(context),
+      cache_path=cache_path,
+      max_chunk_tokens_count=self._max_chunk_tokens_count,
+      max_threads_count=self._max_threads_count,
+      target_language=self._target_language,
+      user_prompt=self._user_prompt,
+      report_progress=report_progress,
+    ):
+      did_touch_end = False
+
+      if spine_file is not None and \
+        translated_count >= len(translated_texts):
+        spine_file.write_texts(translated_texts)
+        spine_file = None
+
+      while spine_file is None:
+        spine_path = next(spine_paths_iter, None)
+        if spine_path is None:
+          did_touch_end = True
+          break
+        spine_file = context.read_spine_file(spine_path)
+        if spine_file.texts_length == 0:
+          spine_file = None
+          continue
+        translated_texts = [""] * spine_file.texts_length
+        translated_count = 0
+
+      translated_texts[translated_count] = translated_text
+      translated_count += 1
+
+      if did_touch_end:
+        break
+    if spine_file and translated_count > 0:
+      spine_file.write_texts(translated_texts)
+
+    context.write_spine_file(spine_path, spine_file)
+
+def _gen_fragments(context: ZipContext):
+  for spine_path in context.search_spine_paths():
+    spine_file = context.read_spine_file(spine_path)
+    for text in spine_file.read_texts():
+      yield Fragment(
+        text=text,
+        start_incision=Incision.IMPOSSIBLE,
+        end_incision=Incision.IMPOSSIBLE,
+      )
+
+def _clean_path(path: Path) -> Path:
+  if path.exists():
+    if path.is_file():
+      path.unlink()
+    elif path.is_dir():
+      rmtree(path, ignore_errors=True)
+  return path

epub_translator/xml/__init__.py

@@ -0,0 +1,3 @@
+from .encoder import encode, encode_friendly
+from .decoder import decode_friendly
+from .utils import clone