epub-translator 0.0.7__py3-none-any.whl → 0.1.1__py3-none-any.whl

epub_translator/llm/increasable.py

@@ -1,35 +1,35 @@
 class Increaser:
-  def __init__(self, value_range: tuple[float, float] | None):
-    self._value_range: tuple[float, float] | None = value_range
-    self._current: float | None = value_range[0] if value_range is not None else None
+    def __init__(self, value_range: tuple[float, float] | None):
+        self._value_range: tuple[float, float] | None = value_range
+        self._current: float | None = value_range[0] if value_range is not None else None
 
-  @property
-  def current(self) -> float | None:
-    return self._current
+    @property
+    def current(self) -> float | None:
+        return self._current
+
+    def increase(self):
+        if self._value_range is not None and self._current is not None:
+            _, end_value = self._value_range
+            self._current = self._current + 0.5 * (end_value - self._current)
 
-  def increase(self):
-    if self._value_range is None:
-      return
-    _, end_value = self._value_range
-    self._current = self._current + 0.5 * (end_value - self._current)
 
 class Increasable:
-  def __init__(self, param: float | tuple[float, float] | None):
-    self._value_range: tuple[float, float] | None = None
+    def __init__(self, param: float | tuple[float, float] | None):
+        self._value_range: tuple[float, float] | None = None
 
-    if isinstance(param, int):
-      param = float(param)
-    if isinstance(param, float):
-      param = (param, param)
-    if isinstance(param, tuple):
-      if len(param) != 2:
-        raise ValueError(f"Expected a tuple of length 2, got {len(param)}")
-      begin, end = param
-      if isinstance(begin, int):
-        begin = float(begin)
-      if isinstance(end, int):
-        end = float(end)
-      self._value_range = (begin, end)
+        if isinstance(param, int):
+            param = float(param)
+        if isinstance(param, float):
+            param = (param, param)
+        if isinstance(param, tuple):
+            if len(param) != 2:
+                raise ValueError(f"Expected a tuple of length 2, got {len(param)}")
+            begin, end = param
+            if isinstance(begin, int):
+                begin = float(begin)
+            if isinstance(end, int):
+                end = float(end)
+            self._value_range = (begin, end)
 
-  def context(self) -> Increaser:
-    return Increaser(self._value_range)
+    def context(self) -> Increaser:
+        return Increaser(self._value_range)

epub_translator/llm/types.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import TypeVar
+
+R = TypeVar("R")
+
+
+@dataclass
+class Message:
+    role: "MessageRole"
+    message: str
+
+
+class MessageRole(Enum):
+    SYSTEM = auto()
+    USER = auto()
+    ASSISTANT = auto()

epub_translator/serial/__init__.py

@@ -0,0 +1,2 @@
+from .segment import ST, S, Segment, T
+from .splitter import split

epub_translator/serial/chunk.py

@@ -0,0 +1,52 @@
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass
+from typing import Generic
+
+from resource_segmentation import Resource, Segment, split
+
+from .segment import ST
+
+_INCISION = 0
+
+
+@dataclass
+class Chunk(Generic[ST]):
+    head_remain_tokens: int
+    tail_remain_tokens: int
+    head: list[ST]
+    body: list[ST]
+    tail: list[ST]
+
+
+def split_into_chunks(segments: Iterable[ST], max_group_tokens: int) -> Generator[Chunk[ST], None, None]:
+    for group in split(
+        max_segment_count=max_group_tokens,
+        gap_rate=0.07,
+        tail_rate=0.5,
+        border_incision=_INCISION,
+        resources=(
+            Resource(
+                count=segment.tokens,
+                start_incision=_INCISION,
+                end_incision=_INCISION,
+                payload=segment,
+            )
+            for segment in segments
+        ),
+    ):
+        yield Chunk(
+            head_remain_tokens=group.head_remain_count,
+            tail_remain_tokens=group.tail_remain_count,
+            head=list(_expand_payloads(group.head)),
+            body=list(_expand_payloads(group.body)),
+            tail=list(_expand_payloads(group.tail)),
+        )
+
+
+def _expand_payloads(target: list[Resource[ST] | Segment[ST]]) -> Generator[ST, None, None]:
+    for item in target:
+        if isinstance(item, Resource):
+            yield item.payload
+        elif isinstance(item, Segment):
+            for resource in item.resources:
+                yield resource.payload

epub_translator/serial/segment.py

@@ -0,0 +1,17 @@
+from typing import Generic, Protocol, Self, TypeVar, runtime_checkable
+
+S = TypeVar("S", covariant=True)
+T = TypeVar("T")
+ST = TypeVar("ST", bound="Segment")
+
+
+@runtime_checkable
+class Segment(Protocol, Generic[S]):
+    @property
+    def tokens(self) -> int: ...
+
+    @property
+    def payload(self) -> S: ...
+
+    def truncate_after_head(self, remain_tokens: int) -> Self: ...
+    def truncate_before_tail(self, remain_tokens: int) -> Self: ...

epub_translator/serial/splitter.py

@@ -0,0 +1,50 @@
+from collections.abc import Callable, Generator, Iterable
+
+from .chunk import split_into_chunks
+from .segment import ST, T
+
+
+def split(
+    segments: Iterable[ST],
+    transform: Callable[[list[ST]], list[T]],
+    max_group_tokens: int,
+) -> Generator[T, None, None]:
+    for group in split_into_chunks(segments, max_group_tokens):
+        head = list(
+            _truncate_extra_content(
+                segments=group.head,
+                remain_left=False,
+                remain_tokens=group.head_remain_tokens,
+            )
+        )
+        tail = list(
+            _truncate_extra_content(
+                segments=group.tail,
+                remain_left=True,
+                remain_tokens=group.tail_remain_tokens,
+            )
+        )
+        transformed = transform(head + group.body + tail)
+
+        if len(tail) > 0:  # 避免 target[N:-0] 切片错误
+            yield from transformed[len(head) : -len(tail)]
+        else:
+            yield from transformed[len(head) :]
+
+
+def _truncate_extra_content(segments: list[ST], remain_left: bool, remain_tokens: int):
+    tokens_list: list[int] = [segment.tokens for segment in segments]
+    segments = list(segments)
+    for tokens in tokens_list if remain_left else reversed(tokens_list):
+        if remain_tokens <= 0:
+            break
+        next_segment = segments.pop(0) if remain_left else segments.pop()
+        if remain_tokens < tokens:
+            if remain_left:
+                next_segment = next_segment.truncate_after_head(remain_tokens)
+            else:
+                next_segment = next_segment.truncate_before_tail(remain_tokens)
+            remain_tokens = 0
+        else:
+            remain_tokens -= tokens
+        yield next_segment

epub_translator/template.py

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