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

epub_translator/xml_translator/submitter.py

@@ -1,69 +1,37 @@
-from collections.abc import Iterable
 from xml.etree.ElementTree import Element
 
-from ..xml import iter_with_stack
-from .text_segment import TextPosition, TextSegment, combine_text_segments
+from ..segment import TextSegment, combine_text_segments
+from ..xml import index_of_parent, iter_with_stack
+from .stream_mapper import InlineSegmentMapping
 
 
-def submit_text_segments(element: Element, text_segments: Iterable[TextSegment]):
-    grouped_map = _group_text_segments(text_segments)
-    flatten_text_segments = dict(_extract_flatten_text_segments(element, grouped_map))
+def submit_text_segments(element: Element, mappings: list[InlineSegmentMapping]) -> Element:
+    grouped_map = _group_text_segments(mappings)
     _append_text_segments(element, grouped_map)
-    _replace_text_segments(element, flatten_text_segments)
+    return element
 
 
-def _group_text_segments(text_segments: Iterable[TextSegment]):
+def _group_text_segments(mappings: list[InlineSegmentMapping]):
     grouped_map: dict[int, list[TextSegment]] = {}
-    for text_segment in text_segments:
-        parent_id = id(text_segment.block_parent)
-        grouped = grouped_map.get(parent_id, None)
-        if grouped is None:
-            grouped_map[parent_id] = grouped = []
-        grouped_map[parent_id].append(text_segment)
-    return grouped_map
-
-
-# 被覆盖的 block 表示一种偶然现象，由于它的子元素会触发 append 操作，若对它也进行 append 操作阅读顺序会混乱
-# 此时只能在它的所有文本后立即接上翻译后的文本
-def _extract_flatten_text_segments(element: Element, grouped_map: dict[int, list[TextSegment]]):
-    override_parent_ids: set[int] = set()
-    for parents, child_element in iter_with_stack(element):
-        if id(child_element) not in grouped_map:
-            continue
-        for parent in parents[:-1]:
-            parent_id = id(parent)
-            if parent_id in grouped_map:
-                override_parent_ids.add(parent_id)
-
-    if id(element) in grouped_map:
-        override_parent_ids.add(id(element))  # root 不会出现在 parents 中需单独添加
-
-    for parent_id in override_parent_ids:
-        yield parent_id, grouped_map.pop(parent_id)
+    for block_element, text_segments in mappings:
+        parent_id = id(block_element)
+        grouped_map[parent_id] = text_segments
+
+    # TODO: 如下是为了清除嵌入文字的 Block，当前版本忽略了嵌入文字的 Block 概念。
+    #       这是书籍中可能出现的一种情况，虽然不多见。
+    #       例如，作为非叶子的块元素，它的子块元素之间会夹杂文本，当前 collect_next_inline_segment 会忽略这些文字：
+    #       <div>
+    #         Some text before.
+    #         <!-- 只有下一行作为叶子节点的块元素内的文字会被处理 -->
+    #         <div>Paragraph 1.</div>
+    #         Some text in between.
+    #       </div>
+    for _, text_segments in mappings:
+        for text_segment in text_segments:
+            for parent_block in text_segment.parent_stack[: text_segment.block_depth - 1]:
+                grouped_map.pop(id(parent_block), None)
 
-
-def _replace_text_segments(element: Element, text_segments: dict[int, list[TextSegment]]):
-    for _, child_element in iter_with_stack(element):
-        tail_text_segments: list[TextSegment] = []
-        for text_segment in text_segments.get(id(child_element), ()):
-            if text_segment.position == TextPosition.TEXT:
-                child_element.text = _append_text(
-                    origin_text=child_element.text,
-                    append_text=text_segment.text,
-                )
-            elif text_segment.position == TextPosition.TAIL:
-                tail_text_segments.append(text_segment)
-
-        tail_text_segments.sort(key=lambda t: t.index)
-        tail_text_segments.reverse()
-        for cc_element in child_element:
-            if not tail_text_segments:
-                break
-            if cc_element.tail is not None:
-                cc_element.tail = _append_text(
-                    origin_text=cc_element.tail,
-                    append_text=tail_text_segments.pop().text,
-                )
+    return grouped_map
 
 
 def _append_text_segments(element: Element, grouped_map: dict[int, list[TextSegment]]):
@@ -74,7 +42,7 @@ def _append_text_segments(element: Element, grouped_map: dict[int, list[TextSegm
         if not grouped:
             continue
         parent = parents[-1]
-        index = _index_of_parent(parents[-1], child_element)
+        index = index_of_parent(parents[-1], child_element)
         combined = next(
             combine_text_segments(
                 segments=(t.strip_block_parents() for t in grouped),
@@ -86,17 +54,3 @@ def _append_text_segments(element: Element, grouped_map: dict[int, list[TextSegm
             parent.insert(index + 1, combined_element)
             combined_element.tail = child_element.tail
             child_element.tail = None
-
-
-def _index_of_parent(parent: Element, checked_element: Element) -> int:
-    for i, child in enumerate(parent):
-        if child == checked_element:
-            return i
-    raise ValueError("Element not found in parent.")
-
-
-def _append_text(origin_text: str | None, append_text: str) -> str:
-    if origin_text is None:
-        return append_text
-    else:
-        return origin_text + append_text

epub_translator/xml_translator/translator.py

@@ -1,15 +1,14 @@
-from collections.abc import Generator, Iterable
+from collections.abc import Callable, Generator, Iterable
 from typing import TypeVar
 from xml.etree.ElementTree import Element
 
-from ..iter_sync import IterSync
 from ..llm import LLM, Message, MessageRole
-from ..xml import encode_friendly
-from .fill import XMLFill
-from .format import ValidationError, _extract_xml_element
-from .group import XMLGroupContext
-from .progressive_locking import ProgressiveLockingValidator
-from .text_segment import TextSegment
+from ..segment import BlockSegment, InlineSegment, TextSegment
+from ..xml import decode_friendly, encode_friendly
+from .callbacks import Callbacks, FillFailedEvent, warp_callbacks
+from .hill_climbing import HillClimbing
+from .stream_mapper import InlineSegmentMapping, XMLStreamMapper
+from .submitter import submit_text_segments
 
 T = TypeVar("T")
 
@@ -17,66 +16,110 @@ T = TypeVar("T")
 class XMLTranslator:
     def __init__(
         self,
-        llm: LLM,
-        group_context: XMLGroupContext,
+        translation_llm: LLM,
+        fill_llm: LLM,
         target_language: str,
         user_prompt: str | None,
         ignore_translated_error: bool,
         max_retries: int,
         max_fill_displaying_errors: int,
+        max_group_tokens: int,
+        cache_seed_content: str | None = None,
     ) -> None:
-        self._llm: LLM = llm
-        self._group_context: XMLGroupContext = group_context
+        self._translation_llm: LLM = translation_llm
+        self._fill_llm: LLM = fill_llm
         self._target_language: str = target_language
         self._user_prompt: str | None = user_prompt
         self._ignore_translated_error: bool = ignore_translated_error
         self._max_retries: int = max_retries
         self._max_fill_displaying_errors: int = max_fill_displaying_errors
+        self._cache_seed_content: str | None = cache_seed_content
+        self._stream_mapper: XMLStreamMapper = XMLStreamMapper(
+            encoding=translation_llm.encoding,
+            max_group_tokens=max_group_tokens,
+        )
 
-    def translate_to_element(self, element: Element) -> Element:
-        for translated, _, _ in self.translate_to_text_segments(((element, None),)):
+    def translate_element(
+        self,
+        element: Element,
+        interrupt_source_text_segments: Callable[[Iterable[TextSegment]], Iterable[TextSegment]] | None = None,
+        interrupt_translated_text_segments: Callable[[Iterable[TextSegment]], Iterable[TextSegment]] | None = None,
+        interrupt_block_element: Callable[[Element], Element] | None = None,
+        on_fill_failed: Callable[[FillFailedEvent], None] | None = None,
+    ) -> Element:
+        for translated in self.translate_elements(
+            elements=((element),),
+            interrupt_source_text_segments=interrupt_source_text_segments,
+            interrupt_translated_text_segments=interrupt_translated_text_segments,
+            interrupt_block_element=interrupt_block_element,
+            on_fill_failed=on_fill_failed,
+        ):
             return translated
-        raise RuntimeError("Translation failed unexpectedly")
 
-    def translate_to_text_segments(
-        self, items: Iterable[tuple[Element, T]]
-    ) -> Generator[tuple[Element, list[TextSegment], T], None, None]:
-        sync: IterSync[tuple[Element, T]] = IterSync()
-        text_segments: list[TextSegment] = []
+        raise RuntimeError("Translation failed unexpectedly")
 
-        for text_segment in self._translate_text_segments(
-            elements=(e for e, _ in sync.iter(items)),
+    def translate_elements(
+        self,
+        elements: Iterable[Element],
+        interrupt_source_text_segments: Callable[[Iterable[TextSegment]], Iterable[TextSegment]] | None = None,
+        interrupt_translated_text_segments: Callable[[Iterable[TextSegment]], Iterable[TextSegment]] | None = None,
+        interrupt_block_element: Callable[[Element], Element] | None = None,
+        on_fill_failed: Callable[[FillFailedEvent], None] | None = None,
+    ) -> Generator[Element, None, None]:
+        callbacks = warp_callbacks(
+            interrupt_source_text_segments=interrupt_source_text_segments,
+            interrupt_translated_text_segments=interrupt_translated_text_segments,
+            interrupt_block_element=interrupt_block_element,
+            on_fill_failed=on_fill_failed,
+        )
+        for element, mappings in self._stream_mapper.map_stream(
+            elements=iter(elements),
+            callbacks=callbacks,
+            map=lambda inline_segments: self._translate_inline_segments(
+                inline_segments=inline_segments,
+                callbacks=callbacks,
+            ),
         ):
-            while True:
-                if sync.tail is None:
-                    break
-                tail_element, _ = sync.tail
-                if id(tail_element) == id(text_segment.root):
-                    break
-                tail_element, payload = sync.take()
-                yield tail_element, text_segments, payload
-                text_segments = []
-            text_segments.append(text_segment)
-
-        while sync.tail is not None:
-            tail_element, payload = sync.take()
-            yield tail_element, text_segments, payload
-            text_segments = []
-
-    def _translate_text_segments(self, elements: Iterable[Element]):
-        for group in self._group_context.split_groups(elements):
-            text_segments = list(group)
-            fill = XMLFill(text_segments)
-            source_text = "".join(self._render_text_segments(text_segments))
-            translated_text = self._translate_text(source_text)
-            self._fill_into_xml(
-                fill=fill,
-                source_text=source_text,
-                translated_text=translated_text,
+            yield submit_text_segments(
+                element=element,
+                mappings=mappings,
             )
-            yield from group.body
+
+    def _translate_inline_segments(
+        self,
+        inline_segments: list[InlineSegment],
+        callbacks: Callbacks,
+    ) -> list[InlineSegmentMapping | None]:
+        hill_climbing = HillClimbing(
+            encoding=self._fill_llm.encoding,
+            max_fill_displaying_errors=self._max_fill_displaying_errors,
+            block_segment=BlockSegment(
+                root_tag="xml",
+                inline_segments=inline_segments,
+            ),
+        )
+        text_segments = (text for inline in inline_segments for text in inline)
+        source_text = "".join(self._render_text_segments(text_segments))
+        translated_text = self._translate_text(source_text)
+
+        self._request_and_submit(
+            hill_climbing=hill_climbing,
+            source_text=source_text,
+            translated_text=translated_text,
+            callbacks=callbacks,
+        )
+        mappings: list[InlineSegmentMapping | None] = []
+        for mapping in hill_climbing.gen_mappings():
+            if mapping:
+                _, text_segments = mapping
+                if not text_segments:
+                    mapping = None
+            mappings.append(mapping)
+
+        return mappings
 
     def _render_text_segments(self, segments: Iterable[TextSegment]):
+        # TODO: 没必要，直接按照新的 inline segment 组织就行了
         iterator = iter(segments)
         segment = next(iterator, None)
         if segment is None:
@@ -92,88 +135,94 @@ class XMLTranslator:
         yield segment.text
 
     def _translate_text(self, text: str) -> str:
-        return self._llm.request(
-            input=[
-                Message(
-                    role=MessageRole.SYSTEM,
-                    message=self._llm.template("translate").render(
-                        target_language=self._target_language,
-                        user_prompt=self._user_prompt,
+        with self._translation_llm.context(cache_seed_content=self._cache_seed_content) as ctx:
+            return ctx.request(
+                input=[
+                    Message(
+                        role=MessageRole.SYSTEM,
+                        message=self._translation_llm.template("translate").render(
+                            target_language=self._target_language,
+                            user_prompt=self._user_prompt,
+                        ),
                     ),
-                ),
-                Message(role=MessageRole.USER, message=text),
-            ]
-        )
+                    Message(role=MessageRole.USER, message=text),
+                ]
+            )
 
-    def _fill_into_xml(self, fill: XMLFill, source_text: str, translated_text: str) -> Element:
+    def _request_and_submit(
+        self,
+        hill_climbing: HillClimbing,
+        source_text: str,
+        translated_text: str,
+        callbacks: Callbacks,
+    ) -> None:
         user_message = (
             f"Source text:\n{source_text}\n\n"
-            f"XML template:\n```XML\n{encode_friendly(fill.request_element)}\n```\n\n"
+            f"XML template:\n```XML\n{encode_friendly(hill_climbing.request_element())}\n```\n\n"
             f"Translated text:\n{translated_text}"
         )
         fixed_messages: list[Message] = [
             Message(
                 role=MessageRole.SYSTEM,
-                message=self._llm.template("fill").render(),
+                message=self._fill_llm.template("fill").render(),
             ),
             Message(
                 role=MessageRole.USER,
                 message=user_message,
             ),
         ]
-
-        validator = ProgressiveLockingValidator()
         conversation_history: list[Message] = []
-        latest_error: ValidationError | None = None
 
-        with self._llm.context() as llm_context:
-            for _ in range(self._max_retries):
-                # Request LLM response
-                response = llm_context.request(
-                    input=fixed_messages + conversation_history,
-                )
+        with self._fill_llm.context(cache_seed_content=self._cache_seed_content) as llm_context:
+            error_message: str | None = None
 
-                try:
-                    # Extract XML from response
-                    validated_element = _extract_xml_element(response)
+            for retry_count in range(self._max_retries):
+                response = llm_context.request(fixed_messages + conversation_history)
+                validated_element = self._extract_xml_element(response)
+                error_message = None
+                if isinstance(validated_element, str):
+                    error_message = validated_element
+                elif isinstance(validated_element, Element):
+                    error_message = hill_climbing.submit(validated_element)
+
+                if error_message is None:
+                    break
 
-                    # Validate with progressive locking
-                    is_complete, error_message, newly_locked = validator.validate_with_locking(
-                        template_ele=fill.request_element,
-                        validated_ele=validated_element,
-                        errors_limit=self._max_fill_displaying_errors,
+                callbacks.on_fill_failed(
+                    FillFailedEvent(
+                        error_message=error_message,
+                        retried_count=retry_count + 1,
+                        over_maximum_retries=False,
                     )
+                )
+                conversation_history = [
+                    Message(role=MessageRole.ASSISTANT, message=response),
+                    Message(role=MessageRole.USER, message=error_message),
+                ]
+            if error_message is not None:
+                callbacks.on_fill_failed(
+                    FillFailedEvent(
+                        error_message=error_message,
+                        retried_count=self._max_retries,
+                        over_maximum_retries=True,
+                    )
+                )
+
+    def _extract_xml_element(self, text: str) -> Element | str:
+        first_xml_element: Element | None = None
+        all_xml_elements: int = 0
 
-                    if is_complete:
-                        # All nodes locked, fill successful
-                        fill._fill_submitted_texts(  # pylint: disable=protected-access
-                            generated_ids_stack=[],
-                            element=validated_element,
-                        )
-                        return validated_element
-
-                    # Not complete yet, construct error message with progress info
-                    progress_msg = f"Progress: {len(validator.locked_ids)} nodes locked"
-                    if newly_locked:
-                        progress_msg += f", {len(newly_locked)} newly locked this round"
-
-                    full_error_message = f"{progress_msg}\n\n{error_message}"
-
-                    conversation_history = [
-                        Message(role=MessageRole.ASSISTANT, message=response),
-                        Message(role=MessageRole.USER, message=full_error_message),
-                    ]
-
-                except ValidationError as error:
-                    # XML extraction or basic validation failed
-                    latest_error = error
-                    conversation_history = [
-                        Message(role=MessageRole.ASSISTANT, message=response),
-                        Message(role=MessageRole.USER, message=str(error)),
-                    ]
-
-            message = f"Failed to get valid XML structure after {self._max_retries} attempts"
-            if latest_error is None:
-                raise ValueError(message)
-            else:
-                raise ValueError(message) from latest_error
+        for xml_element in decode_friendly(text, tags="xml"):
+            if first_xml_element is None:
+                first_xml_element = xml_element
+            all_xml_elements += 1
+
+        if first_xml_element is None:
+            return "No complete <xml>...</xml> block found. Please ensure you have properly closed the XML with </xml> tag."  # noqa: E501
+
+        if all_xml_elements > 1:
+            return (
+                f"Found {all_xml_elements} <xml>...</xml> blocks. "
+                "Please return only one XML block without any examples or explanations."
+            )
+        return first_xml_element