epub-translator 0.0.6__py3-none-any.whl → 0.1.0__py3-none-any.whl

epub_translator/xml_translator/fragmented.py

@@ -0,0 +1,125 @@
+from collections.abc import Generator, Iterable, Iterator
+from enum import Enum, auto
+from xml.etree.ElementTree import Element
+
+from tiktoken import Encoding
+
+from .utils import expand_left_element_texts, expand_right_element_texts, normalize_text_in_element
+
+
+def group_fragmented_elements(
+    encoding: Encoding,
+    elements: Iterable[Element],
+    group_max_tokens: int,
+) -> Generator[list[Element], None, None]:
+    remain_tokens_count: int = group_max_tokens
+    elements_buffer: list[Element] = []
+
+    for element in elements:
+        if remain_tokens_count <= 0:
+            remain_tokens_count = group_max_tokens
+            if elements_buffer:
+                yield elements_buffer
+                elements_buffer = []
+
+        counter = _XMLCounter(encoding, element)
+        cost_tokens_count = counter.advance_tokens(remain_tokens_count)
+        remain_tokens_count -= cost_tokens_count
+        if not counter.can_advance():
+            elements_buffer.append(element)
+            continue
+
+        if elements_buffer:
+            yield elements_buffer
+            elements_buffer = []
+
+        remain_tokens_count = group_max_tokens - cost_tokens_count
+        cost_tokens_count = counter.advance_tokens(remain_tokens_count)
+        if not counter.can_advance():
+            elements_buffer.append(element)
+            remain_tokens_count -= cost_tokens_count
+            continue
+
+        remain_tokens_count = group_max_tokens
+        yield [element]
+
+    if elements_buffer:
+        yield elements_buffer
+
+
+class _TextItemKind(Enum):
+    TEXT = auto()
+    XML_TAG = auto()
+
+
+class _XMLCounter:
+    def __init__(self, encoding: Encoding, root: Element) -> None:
+        self._encoding: Encoding = encoding
+        self._text_iter: Iterator[str] = iter(self._expand_texts(root))
+        self._remain_tokens_count: int = 0
+        self._next_text_buffer: str | None = None
+
+    def can_advance(self) -> bool:
+        if self._remain_tokens_count > 0:
+            return True
+        if self._next_text_buffer is None:
+            self._next_text_buffer = next(self._text_iter, None)
+        return self._next_text_buffer is not None
+
+    def _expand_texts(self, element: Element) -> Generator[str, None, None]:
+        xml_tags_buffer: list[str] = []  # 这类过于碎片化，需拼凑避免 encoding 失效
+        for kind, text in self._expand_text_items(element):
+            if kind == _TextItemKind.XML_TAG:
+                xml_tags_buffer.append(text)
+            elif kind == _TextItemKind.TEXT:
+                if xml_tags_buffer:
+                    yield "".join(xml_tags_buffer)
+                    xml_tags_buffer = []
+                yield text
+        if xml_tags_buffer:
+            yield "".join(xml_tags_buffer)
+
+    def _expand_text_items(self, element: Element) -> Generator[tuple[_TextItemKind, str], None, None]:
+        for text in expand_left_element_texts(element):
+            yield _TextItemKind.XML_TAG, text
+
+        text = normalize_text_in_element(element.text)
+        if text is not None:
+            yield _TextItemKind.TEXT, text
+        for child in element:
+            yield from self._expand_text_items(child)
+            tail = normalize_text_in_element(child.tail)
+            if tail is not None:
+                yield _TextItemKind.TEXT, tail
+
+        for text in expand_right_element_texts(element):
+            yield _TextItemKind.XML_TAG, text
+
+    def advance_tokens(self, max_tokens_count: int) -> int:
+        tokens_count: int = 0
+        while tokens_count < max_tokens_count:
+            if self._remain_tokens_count > 0:
+                will_count_tokens = max_tokens_count - tokens_count
+                if will_count_tokens > self._remain_tokens_count:
+                    tokens_count += self._remain_tokens_count
+                    self._remain_tokens_count = 0
+                else:
+                    tokens_count += will_count_tokens
+                    self._remain_tokens_count -= will_count_tokens
+                if tokens_count >= max_tokens_count:
+                    break
+            next_text = self._next_text()
+            if next_text is None:
+                break
+            self._remain_tokens_count += len(self._encoding.encode(next_text))
+
+        return tokens_count
+
+    def _next_text(self) -> str | None:
+        next_text: str | None = None
+        if self._next_text_buffer is None:
+            next_text = next(self._text_iter, None)
+        else:
+            next_text = self._next_text_buffer
+            self._next_text_buffer = None
+        return next_text

epub_translator/xml_translator/group.py

@@ -0,0 +1,183 @@
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass
+from xml.etree.ElementTree import Element
+
+from resource_segmentation import Resource, Segment, split
+from tiktoken import Encoding
+
+from .fragmented import group_fragmented_elements
+from .text_segment import TextSegment, incision_between, search_text_segments
+
+_BORDER_INCISION = 0
+_ELLIPSIS = "..."
+
+
+@dataclass
+class XMLGroup:
+    head: list[TextSegment]
+    body: list[TextSegment]
+    tail: list[TextSegment]
+
+    def __iter__(self) -> Generator[TextSegment, None, None]:
+        yield from self.head
+        yield from self.body
+        yield from self.tail
+
+
+class XMLGroupContext:
+    def __init__(self, encoding: Encoding, max_group_tokens: int) -> None:
+        self._encoding: Encoding = encoding
+        self._max_group_tokens: int = max_group_tokens
+
+    def split_groups(self, elements: Iterable[Element]) -> Generator[XMLGroup, None, None]:
+        for grouped_elements in group_fragmented_elements(
+            encoding=self._encoding,
+            elements=elements,
+            group_max_tokens=self._max_group_tokens,
+        ):
+            for group in split(
+                resources=self._expand_text_segments(grouped_elements),
+                max_segment_count=self._max_group_tokens,
+                border_incision=_BORDER_INCISION,
+            ):
+                yield XMLGroup(
+                    head=list(
+                        self._truncate_text_segments(
+                            segments=self._expand_text_segments_with_items(group.head),
+                            remain_head=False,
+                            remain_count=group.head_remain_count,
+                        )
+                    ),
+                    body=list(self._expand_text_segments_with_items(group.body)),
+                    tail=list(
+                        self._truncate_text_segments(
+                            segments=self._expand_text_segments_with_items(group.tail),
+                            remain_head=True,
+                            remain_count=group.tail_remain_count,
+                        )
+                    ),
+                )
+
+    def _expand_text_segments(self, elements: Iterable[Element]):
+        for element in elements:
+            yield from self._expand_text_segments_with_element(element)
+
+    def _expand_text_segments_with_element(self, element: Element) -> Generator[Resource[TextSegment], None, None]:
+        generator = search_text_segments(element)
+        segment = next(generator, None)
+        start_incision = _BORDER_INCISION
+        if segment is None:
+            return
+
+        while True:
+            next_segment = next(generator, None)
+            if next_segment is None:
+                break
+            incision1, incision2 = incision_between(
+                segment1=segment,
+                segment2=next_segment,
+            )
+            yield Resource(
+                count=len(self._encoding.encode(segment.xml_text)),
+                start_incision=start_incision,
+                end_incision=incision1,
+                payload=segment,
+            )
+            segment = next_segment
+            start_incision = incision2
+
+        yield Resource(
+            count=len(self._encoding.encode(segment.xml_text)),
+            start_incision=start_incision,
+            end_incision=_BORDER_INCISION,
+            payload=segment,
+        )
+
+    def _expand_text_segments_with_items(self, items: list[Resource[TextSegment] | Segment[TextSegment]]):
+        for item in items:
+            if isinstance(item, Resource):
+                yield item.payload.clone()
+            elif isinstance(item, Segment):
+                for resource in item.resources:
+                    yield resource.payload.clone()
+
+    def _truncate_text_segments(self, segments: Iterable[TextSegment], remain_head: bool, remain_count: int):
+        if remain_head:
+            yield from self._filter_and_remain_segments(
+                segments=segments,
+                remain_head=remain_head,
+                remain_count=remain_count,
+            )
+        else:
+            yield from reversed(
+                list(
+                    self._filter_and_remain_segments(
+                        segments=reversed(list(segments)),
+                        remain_head=remain_head,
+                        remain_count=remain_count,
+                    )
+                )
+            )
+
+    def _filter_and_remain_segments(self, segments: Iterable[TextSegment], remain_head: bool, remain_count: int):
+        for segment in segments:
+            if remain_count <= 0:
+                break
+            raw_xml_text = segment.xml_text
+            tokens = self._encoding.encode(raw_xml_text)
+            tokens_count = len(tokens)
+
+            if tokens_count > remain_count:
+                truncated_segment = self._truncate_text_segment(
+                    segment=segment,
+                    tokens=tokens,
+                    raw_xml_text=raw_xml_text,
+                    remain_head=remain_head,
+                    remain_count=remain_count,
+                )
+                if truncated_segment is not None:
+                    yield truncated_segment
+                break
+
+            yield segment
+            remain_count -= tokens_count
+
+    def _truncate_text_segment(
+        self,
+        segment: TextSegment,
+        tokens: list[int],
+        raw_xml_text: str,
+        remain_head: bool,
+        remain_count: int,
+    ) -> TextSegment | None:
+        # 典型的 xml_text: <tag id="99" data-origin-len="999">Some text</tag>
+        # 如果切割点在前缀 XML 区，则整体舍弃
+        # 如果切割点在后缀 XML 区，则整体保留
+        # 只有刚好切割在正文区，才执行文本截断操作
+        remain_text: str
+        xml_text_head_length = raw_xml_text.find(segment.text)
+
+        if remain_head:
+            remain_xml_text = self._encoding.decode(tokens[:remain_count])  # remain_count cannot be 0 here
+            if len(remain_xml_text) <= xml_text_head_length:
+                return
+            if len(remain_xml_text) >= xml_text_head_length + len(segment.text):
+                return segment
+            remain_text = remain_xml_text[xml_text_head_length:]
+        else:
+            xml_text_tail_length = len(raw_xml_text) - (xml_text_head_length + len(segment.text))
+            remain_xml_text = self._encoding.decode(tokens[-remain_count:])
+            if len(remain_xml_text) <= xml_text_tail_length:
+                return
+            if len(remain_xml_text) >= xml_text_tail_length + len(segment.text):
+                return segment
+            remain_text = remain_xml_text[: len(remain_xml_text) - xml_text_tail_length]
+
+        if not remain_text.strip():
+            return
+
+        if remain_head:
+            segment.text = f"{remain_text} {_ELLIPSIS}"
+        else:
+            segment.text = f"{_ELLIPSIS} {remain_text}"
+        return segment

epub_translator/xml_translator/progressive_locking.py

@@ -0,0 +1,256 @@
+"""渐进式锁定验证器 - 引导 LLM 逐步收敛到正确结果"""
+# pylint: disable=protected-access  # Intentional access to _ValidationContext internals
+
+from xml.etree.ElementTree import Element
+
+from .const import ID_KEY
+from .format import _ValidationContext
+
+
+class ProgressiveLockingValidator:
+    """
+    渐进式锁定验证器：
+    - 验证时收集所有错误（包括已锁定区域）
+    - 识别无错误的子树并锁定
+    - 只报告未锁定区域的错误
+    - 保证收敛：每轮至少锁定一个节点，最多 N 轮完成
+    """
+
+    def __init__(self):
+        self.locked_ids: set[int] = set()
+        self.no_progress_count: int = 0
+        self.lock_history: list[set[int]] = []  # 记录每轮锁定的节点
+
+    def validate_with_locking(
+        self,
+        template_ele: Element,
+        validated_ele: Element,
+        errors_limit: int,
+    ) -> tuple[bool, str | None, set[int]]:
+        """
+        使用渐进式锁定进行验证
+
+        返回：
+        - is_complete: 是否所有节点都已锁定（完成）
+        - error_message: 未锁定区域的错误消息（None 表示无错误）
+        - newly_locked: 本轮新锁定的节点 ID 集合
+        """
+
+        # 1. 执行完整验证（包括已锁定区域）
+        context = _ValidationContext()
+        context.validate(raw_ele=template_ele, validated_ele=validated_ele)
+
+        # 2. 获取所有错误（以路径为 key）
+        all_errors = context._errors
+
+        # 3. 识别可以新锁定的节点
+        newly_locked = self._find_lockable_nodes(template_ele, validated_ele, all_errors)
+
+        # 4. 检测卡住情况并解锁
+        if not newly_locked and self.locked_ids:
+            self.no_progress_count += 1
+            if self.no_progress_count >= 3:
+                # 卡住了，解锁最近的 2 个节点重试
+                self._unlock_recent(count=2)
+                self.no_progress_count = 0
+        else:
+            self.no_progress_count = 0
+
+        # 5. 更新锁定集合
+        self.locked_ids.update(newly_locked)
+        self.lock_history.append(newly_locked.copy())
+
+        # 6. 过滤错误：只保留未锁定区域的错误
+        unlocked_errors = self._filter_unlocked_errors(all_errors)
+
+        # 7. 生成错误消息
+        error_message = self._format_errors(unlocked_errors, errors_limit, template_ele)
+
+        # 8. 检查是否完成
+        total_nodes = self._count_nodes_with_id(template_ele)
+        is_complete = len(self.locked_ids) == total_nodes and error_message is None
+
+        return is_complete, error_message, newly_locked
+
+    def _find_lockable_nodes(
+        self, template_ele: Element, validated_ele: Element, errors: dict[tuple[int, ...], list[str]]
+    ) -> set[int]:
+        """
+        找到可以锁定的节点（该节点及所有后代都无错误，且尚未锁定）
+
+        策略：自底向上，优先锁定叶子节点
+        """
+        lockable = set()
+
+        # 收集所有带 id 的节点，按深度排序（深度优先）
+        nodes_with_depth = []
+        for elem in template_ele.iter():
+            elem_id_str = elem.get(ID_KEY)
+            if elem_id_str is not None:
+                elem_id = int(elem_id_str)
+                if elem_id not in self.locked_ids:
+                    depth = self._get_depth(elem, template_ele)
+                    nodes_with_depth.append((depth, elem_id, elem))
+
+        # 从最深的节点开始检查
+        nodes_with_depth.sort(reverse=True, key=lambda x: x[0])
+
+        for depth, elem_id, elem in nodes_with_depth:
+            # 检查该节点的子树是否完全无错误
+            if self._subtree_is_error_free(elem, template_ele, errors):
+                # 验证在 validated_ele 中也存在对应节点
+                validated_node = self._find_by_id(validated_ele, elem_id)
+                if validated_node is not None:
+                    lockable.add(elem_id)
+
+        return lockable
+
+    def _subtree_is_error_free(
+        self, root: Element, template_root: Element, errors: dict[tuple[int, ...], list[str]]
+    ) -> bool:
+        """检查子树是否完全无错误"""
+
+        # 获取该节点在 template 中的路径
+        root_path = self._get_path_to_node(root, template_root)
+        if root_path is None:
+            return False
+
+        # 检查该路径及其所有后代路径是否有错误
+        for error_path in errors.keys():
+            # 如果错误路径是 root_path 的后代或等于 root_path
+            if self._is_descendant_path(error_path, root_path):
+                return False
+
+        return True
+
+    def _get_path_to_node(self, target: Element, root: Element) -> tuple[int, ...] | None:
+        """获取从 root 到 target 的路径（以 id 序列表示）"""
+
+        def find_path(current: Element, path: list[int]) -> list[int] | None:
+            if current is target:
+                return path
+
+            current_id_str = current.get(ID_KEY)
+            if current_id_str is not None:
+                current_path = path + [int(current_id_str)]
+            else:
+                current_path = path
+
+            for child in current:
+                result = find_path(child, current_path)
+                if result is not None:
+                    return result
+
+            return None
+
+        path = find_path(root, [])
+        return tuple(path) if path is not None else None
+
+    def _is_descendant_path(self, path: tuple[int, ...], ancestor_path: tuple[int, ...]) -> bool:
+        """检查 path 是否是 ancestor_path 的后代或相等"""
+        if len(path) < len(ancestor_path):
+            return False
+        return path[: len(ancestor_path)] == ancestor_path
+
+    def _get_depth(self, elem: Element, root: Element) -> int:
+        """获取元素的深度"""
+        path = self._get_path_to_node(elem, root)
+        return len(path) if path else 0
+
+    def _find_by_id(self, root: Element, target_id: int) -> Element | None:
+        """在树中查找指定 id 的元素"""
+        for elem in root.iter():
+            elem_id_str = elem.get(ID_KEY)
+            if elem_id_str is not None and int(elem_id_str) == target_id:
+                return elem
+        return None
+
+    def _filter_unlocked_errors(self, errors: dict[tuple[int, ...], list[str]]) -> dict[tuple[int, ...], list[str]]:
+        """过滤错误：只保留路径中包含未锁定节点的错误"""
+        unlocked_errors = {}
+
+        for path, error_list in errors.items():
+            # 检查路径中是否有未锁定的节点
+            has_unlocked = any(node_id not in self.locked_ids for node_id in path)
+            if has_unlocked:
+                unlocked_errors[path] = error_list
+
+        return unlocked_errors
+
+    def _format_errors(
+        self, errors: dict[tuple[int, ...], list[str]], limit: int, template_ele: Element
+    ) -> str | None:
+        """格式化错误消息（复用现有逻辑）"""
+        if not errors:
+            return None
+
+        # 使用现有的错误格式化逻辑
+        context = _ValidationContext()
+        context._errors = errors
+
+        # 构造 _tag_text_dict，从 template_ele 中提取真实的标签信息
+        id_to_elem: dict[int, Element] = {}
+        for elem in template_ele.iter():
+            elem_id_str = elem.get(ID_KEY)
+            if elem_id_str is not None:
+                elem_id = int(elem_id_str)
+                id_to_elem[elem_id] = elem
+
+        # 填充 _tag_text_dict
+        for path in errors.keys():
+            for node_id in path:
+                if node_id not in context._tag_text_dict:
+                    elem = id_to_elem.get(node_id)
+                    if elem is not None:
+                        context._tag_text_dict[node_id] = self._str_tag(elem)
+                    else:
+                        context._tag_text_dict[node_id] = f'<tag id="{node_id}">'
+
+        return context.errors(limit=limit)
+
+    def _str_tag(self, ele: Element) -> str:
+        """生成标签的字符串表示（与 format.py 中的逻辑一致）"""
+        ele_id = ele.get(ID_KEY)
+        content: str
+        if ele_id is not None:
+            content = f'<{ele.tag} id="{ele_id}"'
+        else:
+            content = f"<{ele.tag}"
+        if len(ele) > 0:
+            content += f"> ... </{ele.tag}>"
+        else:
+            content += " />"
+        return content
+
+    def _count_nodes_with_id(self, root: Element) -> int:
+        """统计带有 id 属性的节点数量"""
+        count = 0
+        for elem in root.iter():
+            if elem.get(ID_KEY) is not None:
+                count += 1
+        return count
+
+    def _unlock_recent(self, count: int):
+        """解锁最近锁定的 count 个节点"""
+        if not self.lock_history:
+            return
+
+        unlocked_count = 0
+        # 从最近的历史记录开始解锁
+        for i in range(len(self.lock_history) - 1, -1, -1):
+            if unlocked_count >= count:
+                break
+
+            locked_in_round = self.lock_history[i]
+            for node_id in locked_in_round:
+                if unlocked_count >= count:
+                    break
+                if node_id in self.locked_ids:
+                    self.locked_ids.remove(node_id)
+                    unlocked_count += 1
+
+    def get_progress_summary(self, total_nodes: int) -> str:
+        """获取进度摘要"""
+        locked_count = len(self.locked_ids)
+        percentage = (locked_count / total_nodes * 100) if total_nodes > 0 else 0
+        return f"{locked_count}/{total_nodes} nodes locked ({percentage:.1f}%)"

epub_translator/xml_translator/submitter.py

@@ -0,0 +1,102 @@
+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
+
+
+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))
+    _append_text_segments(element, grouped_map)
+    _replace_text_segments(element, flatten_text_segments)
+
+
+def _group_text_segments(text_segments: Iterable[TextSegment]):
+    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)
+
+
+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,
+                )
+
+
+def _append_text_segments(element: Element, grouped_map: dict[int, list[TextSegment]]):
+    for parents, child_element in iter_with_stack(element):
+        if not parents:
+            continue
+        grouped = grouped_map.get(id(child_element))
+        if not grouped:
+            continue
+        parent = parents[-1]
+        index = _index_of_parent(parents[-1], child_element)
+        combined = next(
+            combine_text_segments(
+                segments=(t.strip_block_parents() for t in grouped),
+            ),
+            None,
+        )
+        if combined is not None:
+            combined_element, _ = combined
+            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