File size: 5,167 Bytes
86c402d
744a170
86c402d
 
be03119
744a170
86c402d
 
 
 
744a170
 
 
 
 
86c402d
 
 
744a170
 
86c402d
744a170
 
 
 
 
86c402d
 
744a170
86c402d
744a170
 
 
86c402d
744a170
86c402d
 
 
308de05
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
744a170
 
 
 
 
 
86c402d
744a170
 
86c402d
744a170
 
 
 
 
86c402d
744a170
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
be03119
 
 
 
 
 
 
 
 
 
 
 
744a170
 
 
 
 
 
 
 
 
86c402d
744a170
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
"""
Абстрактный базовый класс для стратегий чанкинга.
"""

import asyncio
import logging
from abc import ABC, abstractmethod

from ntr_fileparser import ParsedDocument

from ..models import DocumentAsEntity, LinkerEntity
from ..repositories import EntityRepository
from .models import Chunk

logger = logging.getLogger(__name__)


class ChunkingStrategy(ABC):
    """Абстрактный класс для стратегий чанкинга."""

    @abstractmethod
    def chunk(
        self,
        document: ParsedDocument,
        doc_entity: DocumentAsEntity,
    ) -> list[LinkerEntity]:
        """
        Разбивает документ на чанки в соответствии со стратегией.

        Args:
            document: ParsedDocument для извлечения текста и структуры.
            doc_entity: Сущность документа-владельца, к которой будут привязаны чанки.

        Returns:
            Список сущностей (чанки)
        """
        raise NotImplementedError("Стратегия чанкинга должна реализовать метод chunk")

    @abstractmethod
    async def chunk_async(
        self,
        document: ParsedDocument,
        doc_entity: DocumentAsEntity,
    ) -> list[LinkerEntity]:
        """
        Асинхронно разбивает документ на чанки в соответствии со стратегией.

        Args:
            document: ParsedDocument для извлечения текста и структуры.
            doc_entity: Сущность документа-владельца, к которой будут привязаны чанки.

        Returns:
            Список сущностей (чанки)
        """
        logger.warning(
            "Асинхронная стратегия чанкинга не реализована, вызывается синхронная"
        )
        return self.chunk(document, doc_entity)

    @classmethod
    def dechunk(
        cls,
        repository: EntityRepository,
        filtered_entities: list[LinkerEntity],
    ) -> str:
        """
        Собирает текст из отфильтрованных чанков к одному документу.

        Args:
            repository: Репозиторий (может понадобиться для получения доп. информации,
                        хотя в текущей реализации не используется).
            filtered_entities: Список отфильтрованных сущностей (чанков),
                               относящихся к одному документу.

        Returns:
            Собранный текст из чанков.
        """
        chunks = [e for e in filtered_entities if isinstance(e, Chunk)]
        chunks.sort(key=lambda x: x.number_in_relation)

        groups: list[list[Chunk]] = []
        for chunk in chunks:
            if len(groups) == 0:
                groups.append([chunk])
                continue

            last_chunk = groups[-1][-1]
            if chunk.number_in_relation == last_chunk.number_in_relation + 1:
                groups[-1].append(chunk)
            else:
                groups.append([chunk])

        result = ""
        previous_last_index = 0
        for group in groups:
            if previous_last_index is not None:
                missing_chunks = group[0].number_in_relation - previous_last_index - 1
                missing_string = f'\n_<...Пропущено {missing_chunks} фрагментов...>_\n'
            else:
                missing_string = '\n_<...>_\n'
            result += missing_string + cls._build_sequenced_chunks(repository, group)
            previous_last_index = group[-1].number_in_relation

        return result.strip()

    @classmethod
    async def dechunk_async(
        cls,
        repository: EntityRepository,
        filtered_entities: list[LinkerEntity],
    ) -> str:
        """
        Асинхронно собирает текст из отфильтрованных чанков к одному документу.
        По умолчанию вызывает синхронную версию.
        """
        return await asyncio.to_thread(cls.dechunk, repository, filtered_entities)

    @classmethod
    def _build_sequenced_chunks(
        cls,
        repository: EntityRepository,
        group: list[Chunk],
    ) -> str:
        """
        Строит текст для последовательных чанков.
        Стоит переопределить в конкретной стратегии, если она предполагает сложную логику
        """
        return " ".join([cls._build_chunk(chunk) for chunk in group])

    @classmethod
    def _build_chunk(cls, chunk: Chunk) -> str:
        """Строит текст для одного чанка."""
        return chunk.text