diff --git a/Dockerfile b/Dockerfile
index cad296ea99cbd2bdfce7853508218a0a30b42764..e762b291ea9cb67e566fd9cda61731eeadede17f 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -30,13 +30,16 @@ RUN python -m pip install \
torch==2.6.0+cu126 \
--index-url https://download.pytorch.org/whl/cu126
+
COPY requirements.txt /app/
RUN python -m pip install -r requirements.txt
-# RUN python -m pip install --ignore-installed elasticsearch==7.11.0 || true
COPY . .
+RUN python -m pip install -e ./lib/parser
+RUN python -m pip install --no-deps -e ./lib/extractor
+# RUN python -m pip install --ignore-installed elasticsearch==7.11.0 || true
-# RUN mkdir -p /data/regulation_datasets /data/documents /data/logs
+RUN mkdir -p /data/regulation_datasets /data/documents /logs
EXPOSE ${PORT}
diff --git a/common/db.py b/common/db.py
index a0048e28c212faea2976c72ec06eaa8d08027f24..0a46ed4493e8b8cf125f3dcf6380ea7a70d2d1d2 100644
--- a/common/db.py
+++ b/common/db.py
@@ -16,13 +16,12 @@ import components.dbo.models.document
import components.dbo.models.log
import components.dbo.models.llm_prompt
import components.dbo.models.llm_config
-
+import components.dbo.models.entity
CONFIG_PATH = os.environ.get('CONFIG_PATH', './config_dev.yaml')
config = Configuration(CONFIG_PATH)
logger = logging.getLogger(__name__)
-print("sql url:", config.common_config.log_sql_path)
engine = create_engine(config.common_config.log_sql_path, connect_args={'check_same_thread': False})
session_factory = sessionmaker(autocommit=False, autoflush=False, bind=engine)
diff --git a/common/dependencies.py b/common/dependencies.py
index 22fe76ba917a65ce33fdca5ce593841840797433..144734ef9f11d0a5fc966b24cddb6b686ca203e0 100644
--- a/common/dependencies.py
+++ b/common/dependencies.py
@@ -1,21 +1,22 @@
import logging
-from logging import Logger
import os
+from logging import Logger
+from typing import Annotated
+
from fastapi import Depends
+from ntr_text_fragmentation import InjectionBuilder
+from sqlalchemy.orm import Session, sessionmaker
from common.configuration import Configuration
+from common.db import session_factory
+from components.dbo.chunk_repository import ChunkRepository
+from components.embedding_extraction import EmbeddingExtractor
from components.llm.common import LlmParams
from components.llm.deepinfra_api import DeepInfraApi
from components.services.dataset import DatasetService
-from components.embedding_extraction import EmbeddingExtractor
-from components.datasets.dispatcher import Dispatcher
from components.services.document import DocumentService
-from components.services.acronym import AcronymService
+from components.services.entity import EntityService
from components.services.llm_config import LLMConfigService
-
-from typing import Annotated
-from sqlalchemy.orm import sessionmaker, Session
-from common.db import session_factory
from components.services.llm_prompt import LlmPromptService
@@ -28,56 +29,76 @@ def get_db() -> sessionmaker:
def get_logger() -> Logger:
- return logging.getLogger(__name__)
+ return logging.getLogger(__name__)
-def get_embedding_extractor(config: Annotated[Configuration, Depends(get_config)]) -> EmbeddingExtractor:
+def get_embedding_extractor(
+ config: Annotated[Configuration, Depends(get_config)],
+) -> EmbeddingExtractor:
return EmbeddingExtractor(
config.db_config.faiss.model_embedding_path,
config.db_config.faiss.device,
)
-def get_dataset_service(
+def get_chunk_repository(db: Annotated[Session, Depends(get_db)]) -> ChunkRepository:
+ return ChunkRepository(db)
+
+
+def get_injection_builder(
+ chunk_repository: Annotated[ChunkRepository, Depends(get_chunk_repository)],
+) -> InjectionBuilder:
+ return InjectionBuilder(chunk_repository)
+
+
+def get_entity_service(
vectorizer: Annotated[EmbeddingExtractor, Depends(get_embedding_extractor)],
+ chunk_repository: Annotated[ChunkRepository, Depends(get_chunk_repository)],
config: Annotated[Configuration, Depends(get_config)],
- db: Annotated[sessionmaker, Depends(get_db)]
-) -> DatasetService:
- return DatasetService(vectorizer, config, db)
+) -> EntityService:
+ """Получение сервиса для работы с сущностями через DI."""
+ return EntityService(vectorizer, chunk_repository, config)
-def get_dispatcher(vectorizer: Annotated[EmbeddingExtractor, Depends(get_embedding_extractor)],
- config: Annotated[Configuration, Depends(get_config)],
- logger: Annotated[Logger, Depends(get_logger)],
- dataset_service: Annotated[DatasetService, Depends(get_dataset_service)]) -> Dispatcher:
- return Dispatcher(vectorizer, config, logger, dataset_service)
-
-def get_acronym_service(db: Annotated[Session, Depends(get_db)]) -> AcronymService:
- return AcronymService(db)
+def get_dataset_service(
+ entity_service: Annotated[EntityService, Depends(get_entity_service)],
+ config: Annotated[Configuration, Depends(get_config)],
+ db: Annotated[sessionmaker, Depends(get_db)],
+) -> DatasetService:
+ """Получение сервиса для работы с датасетами через DI."""
+ return DatasetService(entity_service, config, db)
-def get_document_service(dataset_service: Annotated[DatasetService, Depends(get_dataset_service)],
- config: Annotated[Configuration, Depends(get_config)],
- db: Annotated[sessionmaker, Depends(get_db)]) -> DocumentService:
+def get_document_service(
+ dataset_service: Annotated[DatasetService, Depends(get_dataset_service)],
+ config: Annotated[Configuration, Depends(get_config)],
+ db: Annotated[sessionmaker, Depends(get_db)],
+) -> DocumentService:
return DocumentService(dataset_service, config, db)
def get_llm_config_service(db: Annotated[Session, Depends(get_db)]) -> LLMConfigService:
return LLMConfigService(db)
-def get_llm_service(config: Annotated[Configuration, Depends(get_config)]) -> DeepInfraApi:
-
- llm_params = LlmParams(**{
- "url": config.llm_config.base_url,
- "model": config.llm_config.model,
- "tokenizer": config.llm_config.tokenizer,
- "type": "deepinfra",
- "default": True,
- "predict_params": None, #должны задаваться при каждом запросе
- "api_key": os.environ.get(config.llm_config.api_key_env),
- "context_length": 128000
- })
+
+def get_llm_service(
+ config: Annotated[Configuration, Depends(get_config)],
+) -> DeepInfraApi:
+
+ llm_params = LlmParams(
+ **{
+ "url": config.llm_config.base_url,
+ "model": config.llm_config.model,
+ "tokenizer": config.llm_config.tokenizer,
+ "type": "deepinfra",
+ "default": True,
+ "predict_params": None, # должны задаваться при каждом запросе
+ "api_key": os.environ.get(config.llm_config.api_key_env),
+ "context_length": 128000,
+ }
+ )
return DeepInfraApi(params=llm_params)
+
def get_llm_prompt_service(db: Annotated[Session, Depends(get_db)]) -> LlmPromptService:
- return LlmPromptService(db)
\ No newline at end of file
+ return LlmPromptService(db)
diff --git a/components/dbo/chunk_repository.py b/components/dbo/chunk_repository.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0f748997ea45f525f31625970fb98f0b88d8f90
--- /dev/null
+++ b/components/dbo/chunk_repository.py
@@ -0,0 +1,249 @@
+from uuid import UUID
+
+import numpy as np
+from ntr_text_fragmentation import LinkerEntity
+from ntr_text_fragmentation.integrations import SQLAlchemyEntityRepository
+from sqlalchemy import and_, select
+from sqlalchemy.orm import Session
+
+from components.dbo.models.entity import EntityModel
+
+
+class ChunkRepository(SQLAlchemyEntityRepository):
+ def __init__(self, db: Session):
+ super().__init__(db)
+
+ def _entity_model_class(self):
+ return EntityModel
+
+ def _map_db_entity_to_linker_entity(self, db_entity: EntityModel):
+ """
+ Преобразует сущность из базы данных в LinkerEntity.
+
+ Args:
+ db_entity: Сущность из базы данных
+
+ Returns:
+ LinkerEntity
+ """
+ # Преобразуем строковые ID в UUID
+ entity = LinkerEntity(
+ id=UUID(db_entity.uuid), # Преобразуем строку в UUID
+ name=db_entity.name,
+ text=db_entity.text,
+ type=db_entity.entity_type,
+ in_search_text=db_entity.in_search_text,
+ metadata=db_entity.metadata_json,
+ source_id=UUID(db_entity.source_id) if db_entity.source_id else None, # Преобразуем строку в UUID
+ target_id=UUID(db_entity.target_id) if db_entity.target_id else None, # Преобразуем строку в UUID
+ number_in_relation=db_entity.number_in_relation,
+ )
+ return LinkerEntity.deserialize(entity)
+
+ def add_entities(
+ self,
+ entities: list[LinkerEntity],
+ dataset_id: int,
+ embeddings: dict[str, np.ndarray],
+ ):
+ """
+ Добавляет сущности в базу данных.
+
+ Args:
+ entities: Список сущностей для добавления
+ dataset_id: ID датасета
+ embeddings: Словарь эмбеддингов {entity_id: embedding}
+ """
+ with self.db() as session:
+ for entity in entities:
+ # Преобразуем UUID в строку для хранения в базе
+ entity_id = str(entity.id)
+
+ if entity_id in embeddings:
+ embedding = embeddings[entity_id]
+ else:
+ embedding = None
+
+ session.add(
+ EntityModel(
+ uuid=str(entity.id), # UUID в строку
+ name=entity.name,
+ text=entity.text,
+ entity_type=entity.type,
+ in_search_text=entity.in_search_text,
+ metadata_json=entity.metadata,
+ source_id=str(entity.source_id) if entity.source_id else None, # UUID в строку
+ target_id=str(entity.target_id) if entity.target_id else None, # UUID в строку
+ number_in_relation=entity.number_in_relation,
+ chunk_index=getattr(entity, "chunk_index", None), # Добавляем chunk_index
+ dataset_id=dataset_id,
+ embedding=embedding,
+ )
+ )
+
+ session.commit()
+
+ def get_searching_entities(
+ self,
+ dataset_id: int,
+ ) -> tuple[list[LinkerEntity], list[np.ndarray]]:
+ with self.db() as session:
+ models = (
+ session.query(EntityModel)
+ .filter(EntityModel.in_search_text is not None)
+ .filter(EntityModel.dataset_id == dataset_id)
+ .all()
+ )
+ return (
+ [self._map_db_entity_to_linker_entity(model) for model in models],
+ [model.embedding for model in models],
+ )
+
+ def get_chunks_by_ids(
+ self,
+ chunk_ids: list[str],
+ ) -> list[LinkerEntity]:
+ """
+ Получение чанков по их ID.
+
+ Args:
+ chunk_ids: Список ID чанков
+
+ Returns:
+ Список чанков
+ """
+ # Преобразуем все ID в строки для единообразия
+ str_chunk_ids = [str(chunk_id) for chunk_id in chunk_ids]
+
+ with self.db() as session:
+ models = (
+ session.query(EntityModel)
+ .filter(EntityModel.uuid.in_(str_chunk_ids))
+ .all()
+ )
+ return [self._map_db_entity_to_linker_entity(model) for model in models]
+
+ def get_entities_by_ids(self, entity_ids: list[UUID]) -> list[LinkerEntity]:
+ """
+ Получить сущности по списку идентификаторов.
+
+ Args:
+ entity_ids: Список идентификаторов сущностей
+
+ Returns:
+ Список сущностей, соответствующих указанным идентификаторам
+ """
+ if not entity_ids:
+ return []
+
+ # Преобразуем UUID в строки
+ str_entity_ids = [str(entity_id) for entity_id in entity_ids]
+
+ with self.db() as session:
+ entity_model = self._entity_model_class()
+ db_entities = session.execute(
+ select(entity_model).where(entity_model.uuid.in_(str_entity_ids))
+ ).scalars().all()
+
+ return [self._map_db_entity_to_linker_entity(entity) for entity in db_entities]
+
+ def get_neighboring_chunks(self, chunk_ids: list[UUID], max_distance: int = 1) -> list[LinkerEntity]:
+ """
+ Получить соседние чанки для указанных чанков.
+
+ Args:
+ chunk_ids: Список идентификаторов чанков
+ max_distance: Максимальное расстояние до соседа
+
+ Returns:
+ Список соседних чанков
+ """
+ if not chunk_ids:
+ return []
+
+ # Преобразуем UUID в строки
+ str_chunk_ids = [str(chunk_id) for chunk_id in chunk_ids]
+
+ with self.db() as session:
+ entity_model = self._entity_model_class()
+ result = []
+
+ # Сначала получаем указанные чанки, чтобы узнать их индексы и документы
+ chunks = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.uuid.in_(str_chunk_ids),
+ entity_model.entity_type == "Chunk" # Используем entity_type вместо type
+ )
+ )
+ ).scalars().all()
+
+ if not chunks:
+ return []
+
+ # Находим документы для чанков через связи
+ doc_ids = set()
+ chunk_indices = {}
+
+ for chunk in chunks:
+ chunk_indices[chunk.uuid] = chunk.chunk_index
+
+ # Находим связь от документа к чанку
+ links = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.target_id == chunk.uuid,
+ entity_model.name == "document_to_chunk"
+ )
+ )
+ ).scalars().all()
+
+ for link in links:
+ doc_ids.add(link.source_id)
+
+ if not doc_ids or not any(idx is not None for idx in chunk_indices.values()):
+ return []
+
+ # Для каждого документа находим все его чанки
+ for doc_id in doc_ids:
+ # Находим все связи от документа к чанкам
+ links = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.source_id == doc_id,
+ entity_model.name == "document_to_chunk"
+ )
+ )
+ ).scalars().all()
+
+ doc_chunk_ids = [link.target_id for link in links]
+
+ # Получаем все чанки документа
+ doc_chunks = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.uuid.in_(doc_chunk_ids),
+ entity_model.entity_type == "Chunk" # Используем entity_type вместо type
+ )
+ )
+ ).scalars().all()
+
+ # Для каждого чанка в документе проверяем, является ли он соседом
+ for doc_chunk in doc_chunks:
+ if doc_chunk.uuid in str_chunk_ids:
+ continue
+
+ if doc_chunk.chunk_index is None:
+ continue
+
+ # Проверяем, является ли чанк соседом какого-либо из исходных чанков
+ is_neighbor = False
+ for orig_chunk_id, orig_index in chunk_indices.items():
+ if orig_index is not None and abs(doc_chunk.chunk_index - orig_index) <= max_distance:
+ is_neighbor = True
+ break
+
+ if is_neighbor:
+ result.append(self._map_db_entity_to_linker_entity(doc_chunk))
+
+ return result
diff --git a/components/dbo/models/dataset.py b/components/dbo/models/dataset.py
index 92dfd6fbd1a289bcd0c75b87d54ec770b692dfcb..1fd396f9a72cb1dfabf646985d4a92dbc75a5504 100644
--- a/components/dbo/models/dataset.py
+++ b/components/dbo/models/dataset.py
@@ -23,4 +23,9 @@ class Dataset(Base):
documents: Mapped[list["DatasetDocument"]] = relationship(
"DatasetDocument", back_populates="dataset",
cascade="all, delete-orphan"
- )
\ No newline at end of file
+ )
+
+ entities: Mapped[list["EntityModel"]] = relationship(
+ "EntityModel", back_populates="dataset",
+ cascade="all, delete-orphan"
+ )
diff --git a/components/dbo/models/entity.py b/components/dbo/models/entity.py
new file mode 100644
index 0000000000000000000000000000000000000000..e85ac166f683fcf5ff1213b3f21e2050307ea3cb
--- /dev/null
+++ b/components/dbo/models/entity.py
@@ -0,0 +1,85 @@
+import json
+
+import numpy as np
+from sqlalchemy import ForeignKey, Integer, LargeBinary, String
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+from sqlalchemy.types import TypeDecorator
+
+from components.dbo.models.base import Base
+
+
+class JSONType(TypeDecorator):
+ """Тип для хранения JSON в SQLite."""
+
+ impl = String
+ cache_ok = True
+
+ def process_bind_param(self, value, dialect):
+ """Сохранение dict в JSON строку."""
+ if value is None:
+ return None
+ return json.dumps(value)
+
+ def process_result_value(self, value, dialect):
+ """Загрузка JSON строки в dict."""
+ if value is None:
+ return None
+ return json.loads(value)
+
+
+class EmbeddingType(TypeDecorator):
+ """Тип для хранения эмбеддингов в SQLite."""
+
+ impl = LargeBinary
+ cache_ok = True
+
+ def process_bind_param(self, value, dialect):
+ """Сохранение numpy array в базу."""
+ if value is None:
+ return None
+ # Убеждаемся, что массив двумерный перед сохранением
+ value = np.asarray(value, dtype=np.float32)
+ if value.ndim == 1:
+ value = value.reshape(1, -1)
+ return value.tobytes()
+
+ def process_result_value(self, value, dialect):
+ """Загрузка из базы в numpy array."""
+ if value is None:
+ return None
+ return np.frombuffer(value, dtype=np.float32)
+
+
+class EntityModel(Base):
+ """
+ SQLAlchemy модель для хранения сущностей.
+ """
+
+ __tablename__ = "entity"
+
+ uuid: Mapped[str] = mapped_column(String, unique=True)
+ name: Mapped[str] = mapped_column(String, nullable=False)
+ text: Mapped[str] = mapped_column(String, nullable=False)
+ in_search_text: Mapped[str] = mapped_column(String, nullable=True)
+ entity_type: Mapped[str] = mapped_column(String, nullable=False)
+
+ # Поля для связей (триплетный подход)
+ source_id: Mapped[str] = mapped_column(String, nullable=True)
+ target_id: Mapped[str] = mapped_column(String, nullable=True)
+ number_in_relation: Mapped[int] = mapped_column(Integer, nullable=True)
+
+ # Поле для индекса чанка в документе
+ chunk_index: Mapped[int] = mapped_column(Integer, nullable=True)
+
+ # JSON-поле для хранения метаданных
+ metadata_json: Mapped[dict] = mapped_column(JSONType, nullable=True)
+
+ embedding: Mapped[np.ndarray] = mapped_column(EmbeddingType, nullable=True)
+
+ dataset_id: Mapped[int] = mapped_column(Integer, ForeignKey("dataset.id"), nullable=False)
+
+ dataset: Mapped["Dataset"] = relationship( # type: ignore
+ "Dataset",
+ back_populates="entities",
+ cascade="all",
+ )
diff --git a/components/embedding_extraction.py b/components/embedding_extraction.py
index 50b0582f27f2b07534516ed2bb9efe4d5d34579f..b971b81ac844b844c8f2fee28a853767cd1c7765 100644
--- a/components/embedding_extraction.py
+++ b/components/embedding_extraction.py
@@ -5,10 +5,10 @@ import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
-from transformers import AutoModel, AutoTokenizer, BatchEncoding, XLMRobertaModel
-from transformers.modeling_outputs import (
- BaseModelOutputWithPoolingAndCrossAttentions as EncoderOutput,
-)
+from transformers import (AutoModel, AutoTokenizer, BatchEncoding,
+ XLMRobertaModel)
+from transformers.modeling_outputs import \
+ BaseModelOutputWithPoolingAndCrossAttentions as EncoderOutput
logger = logging.getLogger(__name__)
@@ -41,8 +41,8 @@ class EmbeddingExtractor:
self.device = device
# Инициализация модели
- self.tokenizer = AutoTokenizer.from_pretrained(model_id)
- self.model: XLMRobertaModel = AutoModel.from_pretrained(model_id).to(
+ self.tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=True)
+ self.model: XLMRobertaModel = AutoModel.from_pretrained(model_id, local_files_only=True).to(
self.device
)
self.model.eval()
@@ -122,7 +122,6 @@ class EmbeddingExtractor:
return embedding.cpu().numpy()
- # TODO: В будущем стоит объединить vectorize и query_embed_extraction
def vectorize(
self,
texts: list[str] | str,
@@ -162,7 +161,11 @@ class EmbeddingExtractor:
logger.info('Vectorized all %d batches', len(embeddings))
- return torch.cat(embeddings).numpy()
+ result = torch.cat(embeddings).numpy()
+ # Всегда возвращаем двумерный массив
+ if result.ndim == 1:
+ result = result.reshape(1, -1)
+ return result
@torch.no_grad()
def _vectorize_batch(
diff --git a/components/nmd/faiss_vector_search.py b/components/nmd/faiss_vector_search.py
index 603238ccfab12cfc2359463a8a9a31485c22b214..b2dd50b95642a92b906c79cc00660446c72a725f 100644
--- a/components/nmd/faiss_vector_search.py
+++ b/components/nmd/faiss_vector_search.py
@@ -1,12 +1,10 @@
import logging
-from typing import List
-import numpy as np
-import pandas as pd
+
import faiss
+import numpy as np
-from common.constants import COLUMN_EMBEDDING
-from common.constants import DO_NORMALIZATION
from common.configuration import DataBaseConfiguration
+from common.constants import DO_NORMALIZATION
from components.embedding_extraction import EmbeddingExtractor
logger = logging.getLogger(__name__)
@@ -14,7 +12,10 @@ logger = logging.getLogger(__name__)
class FaissVectorSearch:
def __init__(
- self, model: EmbeddingExtractor, df: pd.DataFrame, config: DataBaseConfiguration
+ self,
+ model: EmbeddingExtractor,
+ ids_to_embeddings: dict[str, np.ndarray],
+ config: DataBaseConfiguration,
):
self.model = model
self.config = config
@@ -23,26 +24,36 @@ class FaissVectorSearch:
self.k_neighbors = config.ranker.k_neighbors
else:
self.k_neighbors = config.search.vector_search.k_neighbors
- self.__create_index(df)
+ self.index_to_id = {i: id_ for i, id_ in enumerate(ids_to_embeddings.keys())}
+ self.__create_index(ids_to_embeddings)
- def __create_index(self, df: pd.DataFrame):
+ def __create_index(self, ids_to_embeddings: dict[str, np.ndarray]):
"""Load the metadata file."""
- if len(df) == 0:
+ if len(ids_to_embeddings) == 0:
self.index = None
return
- df = df.where(pd.notna(df), None)
- embeddings = np.array(df[COLUMN_EMBEDDING].tolist())
+ embeddings = np.array(list(ids_to_embeddings.values()))
dim = embeddings.shape[1]
- self.index = faiss.IndexFlatL2(dim)
+ self.index = faiss.IndexFlatIP(dim)
self.index.add(embeddings)
def search_vectors(self, query: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Поиск векторов в индексе.
+
+ Args:
+ query: Строка, запрос для поиска.
+
+ Returns:
+ tuple[np.ndarray, np.ndarray, np.ndarray]: Кортеж из трех массивов:
+ - np.ndarray: Вектор запроса (1, embedding_size)
+ - np.ndarray: Оценки косинусного сходства (чем больше, тем лучше)
+ - np.ndarray: Идентификаторы найденных векторов
"""
logger.info(f"Searching vectors in index for query: {query}")
if self.index is None:
return (np.array([]), np.array([]), np.array([]))
query_embeds = self.model.query_embed_extraction(query, DO_NORMALIZATION)
- scores, indexes = self.index.search(query_embeds, self.k_neighbors)
- return query_embeds[0], scores[0], indexes[0]
+ similarities, indexes = self.index.search(query_embeds, self.k_neighbors)
+ ids = [self.index_to_id[index] for index in indexes[0]]
+ return query_embeds, similarities[0], np.array(ids)
diff --git a/components/services/dataset.py b/components/services/dataset.py
index ecd13bfe22a703e7543e2d32d13729bba79ea88c..10eb70c7a789f364afe04521774d61f75a86a969 100644
--- a/components/services/dataset.py
+++ b/components/services/dataset.py
@@ -4,33 +4,27 @@ import os
import shutil
import zipfile
from datetime import datetime
-from multiprocessing import Process
from pathlib import Path
-from typing import Optional
-from threading import Lock
import pandas as pd
import torch
from fastapi import BackgroundTasks, HTTPException, UploadFile
+from ntr_fileparser import ParsedDocument, UniversalParser
+from sqlalchemy.orm import Session
from common.common import get_source_format
from common.configuration import Configuration
-from components.embedding_extraction import EmbeddingExtractor
-from components.parser.features.documents_dataset import DocumentsDataset
-from components.parser.pipeline import DatasetCreationPipeline
-from components.parser.xml.structures import ParsedXML
-from components.parser.xml.xml_parser import XMLParser
-from sqlalchemy.orm import Session
-from components.dbo.models.acronym import Acronym
from components.dbo.models.dataset import Dataset
from components.dbo.models.dataset_document import DatasetDocument
from components.dbo.models.document import Document
+from components.services.entity import EntityService
from schemas.dataset import Dataset as DatasetSchema
from schemas.dataset import DatasetExpanded as DatasetExpandedSchema
from schemas.dataset import DatasetProcessing
from schemas.dataset import DocumentsPage as DocumentsPageSchema
from schemas.dataset import SortQueryList
from schemas.document import Document as DocumentSchema
+
logger = logging.getLogger(__name__)
@@ -38,24 +32,31 @@ class DatasetService:
"""
Сервис для работы с датасетами.
"""
-
+
def __init__(
- self,
- vectorizer: EmbeddingExtractor,
+ self,
+ entity_service: EntityService,
config: Configuration,
- db: Session
+ db: Session,
) -> None:
+ """
+ Инициализация сервиса.
+
+ Args:
+ entity_service: Сервис для работы с сущностями
+ config: Конфигурация приложения
+ db: SQLAlchemy сессия
+ """
logger.info("DatasetService initializing")
self.db = db
self.config = config
- self.parser = XMLParser()
- self.vectorizer = vectorizer
+ self.parser = UniversalParser()
+ self.entity_service = entity_service
self.regulations_path = Path(config.db_config.files.regulations_path)
self.documents_path = Path(config.db_config.files.documents_path)
- self.tmp_path= Path(os.environ.get("APP_TMP_PATH", '.'))
+ self.tmp_path = Path(os.environ.get("APP_TMP_PATH", '.'))
logger.info("DatasetService initialized")
-
def get_dataset(
self,
dataset_id: int,
@@ -83,9 +84,6 @@ class DatasetService:
session.query(Document)
.join(DatasetDocument, DatasetDocument.document_id == Document.id)
.filter(DatasetDocument.dataset_id == dataset_id)
- .filter(
- Document.status.in_(['Актуальный', 'Требует актуализации', 'Упразднён'])
- )
.filter(Document.title.like(f'%{search}%'))
)
@@ -98,7 +96,9 @@ class DatasetService:
.join(DatasetDocument, DatasetDocument.document_id == Document.id)
.filter(DatasetDocument.dataset_id == dataset_id)
.filter(
- Document.status.in_(['Актуальный', 'Требует актуализации', 'Упразднён'])
+ Document.status.in_(
+ ['Актуальный', 'Требует актуализации', 'Упразднён']
+ )
)
.filter(Document.title.like(f'%{search}%'))
.count()
@@ -142,7 +142,7 @@ class DatasetService:
name=dataset.name,
isDraft=dataset.is_draft,
isActive=dataset.is_active,
- dateCreated=dataset.date_created
+ dateCreated=dataset.date_created,
)
for dataset in datasets
]
@@ -198,8 +198,10 @@ class DatasetService:
self.raise_if_processing()
with self.db() as session:
- dataset: Dataset = session.query(Dataset).filter(Dataset.id == dataset_id).first()
-
+ dataset: Dataset = (
+ session.query(Dataset).filter(Dataset.id == dataset_id).first()
+ )
+
if not dataset:
raise HTTPException(status_code=404, detail='Dataset not found')
@@ -222,36 +224,42 @@ class DatasetService:
"""
try:
with self.db() as session:
- dataset = session.query(Dataset).filter(Dataset.id == dataset_id).first()
+ dataset = (
+ session.query(Dataset).filter(Dataset.id == dataset_id).first()
+ )
if not dataset:
- raise HTTPException(status_code=404, detail=f"Dataset with id {dataset_id} not found")
-
- active_dataset = session.query(Dataset).filter(Dataset.is_active == True).first()
-
- self.apply_draft(dataset, session)
+ raise HTTPException(
+ status_code=404,
+ detail=f"Dataset with id {dataset_id} not found",
+ )
+
+ active_dataset = (
+ session.query(Dataset).filter(Dataset.is_active == True).first()
+ )
+
+ self.apply_draft(dataset)
dataset.is_draft = False
dataset.is_active = True
if active_dataset:
active_dataset.is_active = False
-
+
session.commit()
except Exception as e:
logger.error(f"Error applying draft: {e}")
raise
-
-
- def activate_dataset(self, dataset_id: int, background_tasks: BackgroundTasks) -> DatasetExpandedSchema:
+
+ def activate_dataset(
+ self, dataset_id: int, background_tasks: BackgroundTasks
+ ) -> DatasetExpandedSchema:
"""
Активировать датасет в фоновой задаче.
"""
-
+
logger.info(f"Activating dataset {dataset_id}")
self.raise_if_processing()
with self.db() as session:
- dataset = (
- session.query(Dataset).filter(Dataset.id == dataset_id).first()
- )
+ dataset = session.query(Dataset).filter(Dataset.id == dataset_id).first()
active_dataset = session.query(Dataset).filter(Dataset.is_active).first()
if not dataset:
raise HTTPException(status_code=404, detail='Dataset not found')
@@ -329,7 +337,7 @@ class DatasetService:
dataset = self.create_dataset_from_directory(
is_default=False,
- directory_with_xmls=file_location.parent,
+ directory_with_documents=file_location.parent,
directory_with_ready_dataset=None,
)
@@ -341,10 +349,12 @@ class DatasetService:
def apply_draft(
self,
dataset: Dataset,
- session,
) -> None:
"""
Сохранить черновик как полноценный датасет.
+
+ Args:
+ dataset: Датасет для применения
"""
torch.set_num_threads(1)
logger.info(f"Applying draft dataset {dataset.id}")
@@ -363,9 +373,7 @@ class DatasetService:
if current % log_step != 0:
return
if (total > 10) and (current % (total // 10) == 0):
- logger.info(
- f"Processing dataset {dataset.id}: {current}/{total}"
- )
+ logger.info(f"Processing dataset {dataset.id}: {current}/{total}")
with open(TMP_PATH, 'w', encoding='utf-8') as f:
json.dump(
{
@@ -381,34 +389,25 @@ class DatasetService:
document_ids = [
doc_dataset_link.document_id for doc_dataset_link in dataset.documents
]
- document_formats = [
- doc_dataset_link.document.source_format
- for doc_dataset_link in dataset.documents
- ]
-
- prepared_abbreviations = (
- session.query(Acronym).filter(Acronym.document_id.in_(document_ids)).all()
- )
-
- pipeline = DatasetCreationPipeline(
- dataset_id=dataset.id,
- vectorizer=self.vectorizer,
- prepared_abbreviations=prepared_abbreviations,
- document_ids=document_ids,
- document_formats=document_formats,
- datasets_path=self.regulations_path,
- documents_path=self.documents_path,
- save_intermediate_files=True,
- )
- progress_callback(0, 1000)
-
- try:
- pipeline.run(progress_callback)
- except Exception as e:
- logger.error(f"Error running pipeline: {e}")
- raise HTTPException(status_code=500, detail=str(e))
- finally:
- TMP_PATH.unlink()
+
+ for document_id in document_ids:
+ path = self.documents_path / f'{document_id}.DOCX'
+ parsed = self.parser.parse_by_path(str(path))
+ if parsed is None:
+ logger.warning(f"Failed to parse document {document_id}")
+ continue
+
+ # Используем EntityService для обработки документа с callback
+ self.entity_service.process_document(
+ parsed,
+ dataset.id,
+ progress_callback=progress_callback,
+ words_per_chunk=50,
+ overlap_words=25,
+ respect_sentence_boundaries=True,
+ )
+
+ TMP_PATH.unlink()
def raise_if_processing(self) -> None:
"""
@@ -423,7 +422,7 @@ class DatasetService:
def create_dataset_from_directory(
self,
is_default: bool,
- directory_with_xmls: Path,
+ directory_with_documents: Path,
directory_with_ready_dataset: Path | None = None,
) -> Dataset:
"""
@@ -438,7 +437,7 @@ class DatasetService:
Dataset: Созданный датасет.
"""
logger.info(
- f"Creating {'default' if is_default else 'new'} dataset from directory {directory_with_xmls}"
+ f"Creating {'default' if is_default else 'new'} dataset from directory {directory_with_documents}"
)
with self.db() as session:
documents = []
@@ -453,9 +452,9 @@ class DatasetService:
)
session.add(dataset)
- for subpath in self._get_recursive_dirlist(directory_with_xmls):
+ for subpath in self._get_recursive_dirlist(directory_with_documents):
document, relation = self._create_document(
- directory_with_xmls, subpath, dataset
+ directory_with_documents, subpath, dataset
)
if document is None:
continue
@@ -484,7 +483,8 @@ class DatasetService:
old_filename = document.filename
new_filename = '{}.{}'.format(document.id, document.source_format)
shutil.copy(
- directory_with_xmls / old_filename, self.documents_path / new_filename
+ directory_with_documents / old_filename,
+ self.documents_path / new_filename,
)
document.filename = new_filename
@@ -495,16 +495,8 @@ class DatasetService:
dataset_id = dataset.id
-
logger.info(f"Dataset {dataset_id} created")
- df = self.dataset_to_pandas(dataset_id)
-
- (self.regulations_path / str(dataset_id)).mkdir(parents=True, exist_ok=True)
- df.to_csv(
- self.regulations_path / str(dataset_id) / 'documents.csv', index=False
- )
-
return dataset
def create_empty_dataset(self, is_default: bool) -> Dataset:
@@ -526,20 +518,6 @@ class DatasetService:
session.commit()
session.refresh(dataset)
- self.documents_path.mkdir(exist_ok=True)
-
- dataset_id = dataset.id
-
-
- folder = self.regulations_path / str(dataset_id)
- folder.mkdir(parents=True, exist_ok=True)
-
- pickle_creator = DocumentsDataset([])
- pickle_creator.to_pickle(folder / 'dataset.pkl')
-
- df = self.dataset_to_pandas(dataset_id)
- df.to_csv(folder / 'documents.csv', index=False)
-
return dataset
@staticmethod
@@ -553,10 +531,10 @@ class DatasetService:
Returns:
list[Path]: Список путей к xml-файлам относительно path.
"""
- xml_files = set() #set для отбрасывания неуникальных путей
+ xml_files = set() # set для отбрасывания неуникальных путей
for ext in ('*.xml', '*.XML', '*.docx', '*.DOCX'):
xml_files.update(path.glob(f'**/{ext}'))
-
+
return [p.relative_to(path) for p in xml_files]
def _create_document(
@@ -580,19 +558,19 @@ class DatasetService:
try:
source_format = get_source_format(str(subpath))
- parsed_xml: ParsedXML | None = self.parser.parse(
- documents_path / subpath, include_content=False
+ parsed: ParsedDocument | None = self.parser.parse_by_path(
+ str(documents_path / subpath)
)
- if not parsed_xml:
+ if not parsed:
logger.warning(f"Failed to parse file: {subpath}")
return None, None
document = Document(
filename=str(subpath),
- title=parsed_xml.name,
- status=parsed_xml.status,
- owner=parsed_xml.owner,
+ title=parsed.name,
+ status=parsed.meta.status,
+ owner=parsed.meta.owner,
source_format=source_format,
)
relation = DatasetDocument(
@@ -606,36 +584,6 @@ class DatasetService:
logger.error(f"Error creating document from {subpath}: {e}")
return None, None
- def dataset_to_pandas(self, dataset_id: int) -> pd.DataFrame:
- """
- Преобразовать датасет в pandas DataFrame.
- """
- with self.db() as session:
- links = (
- session.query(DatasetDocument)
- .filter(DatasetDocument.dataset_id == dataset_id)
- .all()
- )
- documents = (
- session.query(Document)
- .filter(Document.id.in_([link.document_id for link in links]))
- .all()
- )
-
- return pd.DataFrame(
- [
- {
- 'id': document.id,
- 'filename': document.filename,
- 'title': document.title,
- 'status': document.status,
- 'owner': document.owner,
- }
- for document in documents
- ],
- columns=['id', 'filename', 'title', 'status', 'owner'],
- )
-
def get_current_dataset(self) -> Dataset | None:
with self.db() as session:
print(session)
diff --git a/components/services/document.py b/components/services/document.py
index b23240ab496242401c5ce657372078c39d2038c9..2662ffbd94e0901e0270509f40184403d2aac330 100644
--- a/components/services/document.py
+++ b/components/services/document.py
@@ -4,19 +4,18 @@ import shutil
from pathlib import Path
from fastapi import HTTPException, UploadFile
+from ntr_fileparser import UniversalParser
from sqlalchemy.orm import Session
from common.common import get_source_format
from common.configuration import Configuration
from common.constants import PROCESSING_FORMATS
-from components.parser.xml.xml_parser import XMLParser
from components.dbo.models.dataset import Dataset
from components.dbo.models.dataset_document import DatasetDocument
from components.dbo.models.document import Document
from schemas.document import Document as DocumentSchema
from schemas.document import DocumentDownload
from components.services.dataset import DatasetService
-
logger = logging.getLogger(__name__)
@@ -34,7 +33,7 @@ class DocumentService:
logger.info("Initializing DocumentService")
self.db = db
self.dataset_service = dataset_service
- self.xml_parser = XMLParser()
+ self.parser = UniversalParser()
self.documents_path = Path(config.db_config.files.documents_path)
def get_document(
@@ -101,10 +100,10 @@ class DocumentService:
logger.info(f"Source format: {source_format}")
try:
- parsed = self.xml_parser.parse(file_location, include_content=False)
+ parsed = self.parser.parse_by_path(str(file_location))
except Exception:
raise HTTPException(
- status_code=400, detail="Invalid XML file, service can't parse it"
+ status_code=400, detail="Invalid file, service can't parse it"
)
with self.db() as session:
@@ -118,9 +117,10 @@ class DocumentService:
raise HTTPException(status_code=403, detail='Dataset is not draft')
document = Document(
+ filename=file.filename,
title=parsed.name,
- owner=parsed.owner,
- status=parsed.status,
+ owner=parsed.meta.owner,
+ status=parsed.meta.status,
source_format=source_format,
)
@@ -129,21 +129,21 @@ class DocumentService:
session.add(document)
session.flush()
- logger.info(f"Document ID: {document.document_id}")
+ logger.info(f"Document ID: {document.id}")
link = DatasetDocument(
dataset_id=dataset_id,
- document_id=document.document_id,
+ document_id=document.id,
)
session.add(link)
if source_format in PROCESSING_FORMATS:
logger.info(
- f"Moving file to: {self.documents_path / f'{document.document_id}.{source_format}'}"
+ f"Moving file to: {self.documents_path / f'{document.id}.{source_format}'}"
)
shutil.move(
file_location,
- self.documents_path / f'{document.document_id}.{source_format}',
+ self.documents_path / f'{document.id}.{source_format}',
)
else:
logger.error(f"Unknown source format: {source_format}")
@@ -156,7 +156,7 @@ class DocumentService:
session.refresh(document)
result = DocumentSchema(
- id=document.document_id,
+ id=document.id,
name=document.title,
owner=document.owner,
status=document.status,
diff --git a/components/services/entity.py b/components/services/entity.py
new file mode 100644
index 0000000000000000000000000000000000000000..02c485d1eb1e57ed4e12fb353aa7ce76032b5cc5
--- /dev/null
+++ b/components/services/entity.py
@@ -0,0 +1,210 @@
+import logging
+from typing import Callable, Optional
+from uuid import UUID
+
+import numpy as np
+from ntr_fileparser import ParsedDocument
+from ntr_text_fragmentation import Destructurer, InjectionBuilder, LinkerEntity
+
+from common.configuration import Configuration
+from components.dbo.chunk_repository import ChunkRepository
+from components.embedding_extraction import EmbeddingExtractor
+from components.nmd.faiss_vector_search import FaissVectorSearch
+
+logger = logging.getLogger(__name__)
+
+
+class EntityService:
+ """
+ Сервис для работы с сущностями.
+ Объединяет функциональность chunk_repository, destructurer, injection_builder и faiss_vector_search.
+ """
+
+ def __init__(
+ self,
+ vectorizer: EmbeddingExtractor,
+ chunk_repository: ChunkRepository,
+ config: Configuration,
+ ) -> None:
+ """
+ Инициализация сервиса.
+
+ Args:
+ vectorizer: Модель для извлечения эмбеддингов
+ chunk_repository: Репозиторий для работы с чанками
+ config: Конфигурация приложения
+ """
+ self.vectorizer = vectorizer
+ self.config = config
+ self.chunk_repository = chunk_repository
+ self.faiss_search = None # Инициализируется при необходимости
+ self.current_dataset_id = None # Текущий dataset_id
+
+ def _ensure_faiss_initialized(self, dataset_id: int) -> None:
+ """
+ Проверяет и при необходимости инициализирует или обновляет FAISS индекс.
+
+ Args:
+ dataset_id: ID датасета для инициализации
+ """
+ # Если индекс не инициализирован или датасет изменился
+ if self.faiss_search is None or self.current_dataset_id != dataset_id:
+ logger.info(f'Initializing FAISS for dataset {dataset_id}')
+ entities, embeddings = self.chunk_repository.get_searching_entities(dataset_id)
+ if entities:
+ # Создаем словарь только из не-None эмбеддингов
+ embeddings_dict = {
+ str(entity.id): embedding # Преобразуем UUID в строку для ключа
+ for entity, embedding in zip(entities, embeddings)
+ if embedding is not None
+ }
+ if embeddings_dict: # Проверяем, что есть хотя бы один эмбеддинг
+ self.faiss_search = FaissVectorSearch(
+ self.vectorizer,
+ embeddings_dict,
+ self.config.db_config,
+ )
+ self.current_dataset_id = dataset_id
+ logger.info(f'FAISS initialized for dataset {dataset_id} with {len(embeddings_dict)} embeddings')
+ else:
+ logger.warning(f'No valid embeddings found for dataset {dataset_id}')
+ self.faiss_search = None
+ self.current_dataset_id = None
+ else:
+ logger.warning(f'No entities found for dataset {dataset_id}')
+ self.faiss_search = None
+ self.current_dataset_id = None
+
+ def process_document(
+ self,
+ document: ParsedDocument,
+ dataset_id: int,
+ progress_callback: Optional[Callable] = None,
+ **destructurer_kwargs,
+ ) -> None:
+ """
+ Обработка документа: разбиение на чанки и сохранение в базу.
+
+ Args:
+ document: Документ для обработки
+ dataset_id: ID датасета
+ progress_callback: Функция для отслеживания прогресса
+ **destructurer_kwargs: Дополнительные параметры для Destructurer
+ """
+ logger.info(f"Processing document {document.name} for dataset {dataset_id}")
+
+ # Создаем деструктуризатор с параметрами по умолчанию
+ destructurer = Destructurer(
+ document,
+ strategy_name="fixed_size",
+ process_tables=True,
+ **{
+ "words_per_chunk": 50,
+ "overlap_words": 25,
+ "respect_sentence_boundaries": True,
+ **destructurer_kwargs,
+ }
+ )
+
+ # Получаем сущности
+ entities = destructurer.destructure()
+
+ # Фильтруем сущности для поиска
+ filtering_entities = [entity for entity in entities if entity.in_search_text is not None]
+ filtering_texts = [entity.in_search_text for entity in filtering_entities]
+
+ # Получаем эмбеддинги с поддержкой callback
+ embeddings = self.vectorizer.vectorize(filtering_texts, progress_callback)
+ embeddings_dict = {
+ str(entity.id): embedding # Преобразуем UUID в строку для ключа
+ for entity, embedding in zip(filtering_entities, embeddings)
+ }
+
+ # Сохраняем в базу
+ self.chunk_repository.add_entities(entities, dataset_id, embeddings_dict)
+
+ # Переинициализируем FAISS индекс, если это текущий датасет
+ if self.current_dataset_id == dataset_id:
+ self._ensure_faiss_initialized(dataset_id)
+
+ logger.info(f"Added {len(entities)} entities to dataset {dataset_id}")
+
+ def build_text(
+ self,
+ entities: list[LinkerEntity],
+ chunk_scores: Optional[list[float]] = None,
+ include_tables: bool = True,
+ max_documents: Optional[int] = None,
+ ) -> str:
+ """
+ Сборка текста из сущностей.
+
+ Args:
+ entities: Список сущностей
+ chunk_scores: Список весов чанков
+ include_tables: Флаг включения таблиц
+ max_documents: Максимальное количество документов
+
+ Returns:
+ Собранный текст
+ """
+ logger.info(f"Building text for {len(entities)} entities")
+ if chunk_scores is not None:
+ chunk_scores = {entity.id: score for entity, score in zip(entities, chunk_scores)}
+ builder = InjectionBuilder(self.chunk_repository)
+ return builder.build(
+ [entity.id for entity in entities], # Передаем UUID напрямую
+ chunk_scores=chunk_scores,
+ include_tables=include_tables,
+ max_documents=max_documents,
+ )
+
+ def search_similar(
+ self,
+ query: str,
+ dataset_id: int,
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+ """
+ Поиск похожих сущностей.
+
+ Args:
+ query: Текст запроса
+ dataset_id: ID датасета
+
+ Returns:
+ tuple[np.ndarray, np.ndarray, np.ndarray]:
+ - Вектор запроса
+ - Оценки сходства
+ - Идентификаторы найденных сущностей
+ """
+ # Убеждаемся, что FAISS инициализирован для текущего датасета
+ self._ensure_faiss_initialized(dataset_id)
+
+ if self.faiss_search is None:
+ return np.array([]), np.array([]), np.array([])
+
+ # Выполняем поиск
+ return self.faiss_search.search_vectors(query)
+
+ def add_neighboring_chunks(
+ self,
+ entities: list[LinkerEntity],
+ max_distance: int = 1,
+ ) -> list[LinkerEntity]:
+ """
+ Добавление соседних чанков.
+
+ Args:
+ entities: Список сущностей
+ max_distance: Максимальное расстояние для поиска соседей
+
+ Returns:
+ Расширенный список сущностей
+ """
+ # Убедимся, что все ID представлены в UUID формате
+ for entity in entities:
+ if not isinstance(entity.id, UUID):
+ entity.id = UUID(str(entity.id))
+
+ builder = InjectionBuilder(self.chunk_repository)
+ return builder.add_neighboring_chunks(entities, max_distance)
\ No newline at end of file
diff --git a/lib/extractor/.cursor/rules/project-description.mdc b/lib/extractor/.cursor/rules/project-description.mdc
new file mode 100644
index 0000000000000000000000000000000000000000..54ec9a353584cf2830ab2e9f16ff2aa736e15d51
--- /dev/null
+++ b/lib/extractor/.cursor/rules/project-description.mdc
@@ -0,0 +1,86 @@
+---
+description:
+globs:
+alwaysApply: true
+---
+
+# Project description
+
+Данный проект представляет собой библиотеку, предоставляющую возможности для чанкинга и сборки
+инъекций в промпт LLM для дальнейшего использования в RAG-системах. Основная логика описана в README.md и в architectures, если они не устарели. Ядро системы представляют классы LinkerEntity, Destructurer, EntityRepository, InjectionBuilder, ChunkingStrategy.
+
+- LinkerEntity – основная сущность, от которой затем наследуются Chunk и DocumentAsEntity. Реализует триплетный подход, при котором один и тот же класс задаёт и сущности, и связи, и при этом сущности-ассоциации реализуются одним экземпляром, а не множеством.
+- Destructurer – реализует логику разбиения документа на множество LinkerEntity, во многом делегируя работу различным ChunkingStrategy (но не всю).
+- EntityRepository – интерфейс. Предполагается, что после извлечения всех сущностей посредством Destructurer пользователь библиотеки сохранит все свои сущности некоторым произвольным образом, например, в csv-файл или PostgreSQL. Библиотека не знает, как работать с пользовательскими хранилищами данных, поэтому пользователь должен сам написать реализацию EntityRepository для своего решения, и предоставить её в InjectionBuilder
+- InjectionBuilder – сборщик промпт-инъекции. Принимает на вход отфильтрованный и (в отдельных случаях) оценённый некоторым скором набор сущностей, сортирует их, распределяет по документам и собирает всё в единый текст, пользуясь EntityRepository, чтобы достать связанные полезные сущности
+
+Данная библиотека ориентируется на ParsedDocument из библиотеки ntr_fileparser, структура которого примерно соответствует следующему:
+
+@dataclass
+class ParsedDocument(ParsedStructure):
+ """
+ Документ, полученный в результате парсинга.
+ """
+ name: str = ""
+ type: str = ""
+ meta: ParsedMeta = field(default_factory=ParsedMeta)
+ paragraphs: list[ParsedTextBlock] = field(default_factory=list)
+ tables: list[ParsedTable] = field(default_factory=list)
+ images: list[ParsedImage] = field(default_factory=list)
+ formulas: list[ParsedFormula] = field(default_factory=list)
+
+ def to_string() -> str:
+ ...
+
+ def to_dict() -> dict:
+ ...
+
+
+@dataclass
+class ParsedTextBlock(DocumentElement):
+ """
+ Текстовый блок документа.
+ """
+
+ text: str = ""
+ style: TextStyle = field(default_factory=TextStyle)
+ anchors: list[str] = field(default_factory=list) # Список идентификаторов якорей (закладок)
+ links: list[str] = field(default_factory=list) # Список идентификаторов ссылок
+
+ # Технические метаданные о блоке
+ metadata: list[dict[str, Any]] = field(default_factory=list) # Для хранения технической информации
+
+ # Примечания и сноски к тексту
+ footnotes: list[dict[str, Any]] = field(default_factory=list) # Для хранения сносок
+
+ title_of_table: int | None = None
+
+ def to_string() -> str:
+ ...
+
+ def to_dict() -> dict:
+ ...
+
+
+@dataclass
+class ParsedTable(DocumentElement):
+ """
+ Таблица из документа.
+ """
+
+ title: str | None = None
+ note: str | None = None
+ classified_tags: list[TableTag] = field(default_factory=list)
+ index: list[str] = field(default_factory=list)
+ headers: list[ParsedRow] = field(default_factory=list)
+ subtables: list[ParsedSubtable] = field(default_factory=list)
+ table_style: dict[str, Any] = field(default_factory=dict)
+ title_index_in_paragraphs: int | None = None
+
+ def to_string() -> str:
+ ...
+
+ def to_dict() -> dict:
+ ...
+
+(Дальнейшую информацию о вложенных классах ты можешь уточнить у пользователя, если это будет нужно)
\ No newline at end of file
diff --git a/lib/extractor/.gitignore b/lib/extractor/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..e64e3e5b27654cb393c4678f7e4f4568874c0278
--- /dev/null
+++ b/lib/extractor/.gitignore
@@ -0,0 +1,11 @@
+use_it/*
+test_output/
+test_input/
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+*.pyw
+*.pyz
+
+*.egg-info/
\ No newline at end of file
diff --git a/lib/extractor/README.md b/lib/extractor/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..4c46b4c1cac064613f6e949ad0165a6ebf507964
--- /dev/null
+++ b/lib/extractor/README.md
@@ -0,0 +1,60 @@
+# Библиотека извлечения и сборки документов
+
+Библиотека для извлечения структурированной информации из документов и их последующей сборки.
+
+## Основные компоненты
+
+- **Destructurer**: Разбивает документ на чанки и связи между ними, а также извлекает дополнительные сущности
+- **Builder**: Собирает документ из чанков и связей
+- **Entity**: Базовый класс для всех сущностей (Document, Chunk, Acronym и т.д.)
+- **Link**: Класс для представления связей между сущностями
+- **ChunkingStrategy**: Интерфейс для различных стратегий чанкинга
+- **TablesProcessor**: Процессор для извлечения таблиц из документа
+
+## Установка
+
+```bash
+pip install -e .
+```
+
+## Использование
+
+```python
+from ntr_text_fragmentation.core import Destructurer, Builder
+from ntr_fileparser import ParsedDocument
+
+# Пример использования Destructurer с обработкой таблиц
+document = ParsedDocument(...)
+destructurer = Destructurer(
+ document=document,
+ strategy_name="fixed_size",
+ process_tables=True
+)
+entities = destructurer.destructure()
+
+# Пример использования Builder
+builder = Builder(document)
+builder.configure({"chunking_strategy": "fixed_size"})
+reconstructed_document = builder.build()
+```
+
+## Модули
+
+### Core
+Основные классы для работы с документами:
+- **Destructurer**: Разбивает документ на чанки и другие сущности
+- **Builder**: Собирает документ из чанков и связей
+
+### Chunking
+Различные стратегии разбиения документа на чанки:
+- **FixedSizeChunkingStrategy**: Разбиение на чанки фиксированного размера
+
+### Additors
+Дополнительные обработчики для извлечения сущностей:
+- **TablesProcessor**: Извлекает таблицы из документа и создает для них сущности
+
+### Models
+Модели данных для представления сущностей и связей:
+- **LinkerEntity**: Базовый класс для всех сущностей и связей
+- **DocumentAsEntity**: Представление документа как сущности
+- **TableEntity**: Представление таблицы как сущности
\ No newline at end of file
diff --git a/lib/extractor/docs/architecture.puml b/lib/extractor/docs/architecture.puml
new file mode 100644
index 0000000000000000000000000000000000000000..a5a56f99229113bad4cdee36c7ff97f54cd294ce
--- /dev/null
+++ b/lib/extractor/docs/architecture.puml
@@ -0,0 +1,149 @@
+@startuml "NTR Text Fragmentation Architecture"
+
+' Использование CSS-стилей вместо skinparams
+
+
+' Легенда
+legend
+ Легенда
+
+ | Цвет | Описание |
+ | Зеленый | Модели данных |
+ | Голубой | Стратегии чанкинга |
+ | Красный | Основные компоненты |
+endlegend
+
+' Разделение на пакеты
+
+package "models" {
+ class LinkerEntity <> {
+ + id: UUID
+ + name: str
+ + text: str
+ + in_search_text: str | None
+ + metadata: dict
+ + source_id: UUID | None
+ + target_id: UUID | None
+ + number_in_relation: int | None
+ + type: str
+ + serialize(): LinkerEntity
+ + {abstract} deserialize(data: LinkerEntity): Self
+ }
+
+ class Chunk <> extends LinkerEntity {
+ + chunk_index: int | None
+ }
+
+ class DocumentAsEntity <> extends LinkerEntity {
+ }
+
+ note right of LinkerEntity
+ Базовая сущность для всех элементов системы.
+ in_search_text определяет текст, используемый
+ при поиске, если None - данная сущность не должна попасть
+ в поиск и используется только для вспомогательных целей.
+ end note
+}
+
+package "chunking_strategies" as chunking_strategies {
+ abstract class ChunkingStrategy <> {
+ + {abstract} chunk(document: ParsedDocument, doc_entity: DocumentAsEntity): list[LinkerEntity]
+ + dechunk(entities: list[LinkerEntity], links: list[LinkerEntity]): str
+ }
+
+ package "specific_strategies" {
+ class FixedSizeChunkingStrategy <> extends chunking_strategies.ChunkingStrategy {
+ + chunk(document: ParsedDocument, doc_entity: DocumentAsEntity): list[LinkerEntity]
+ + dechunk(entities: list[LinkerEntity], links: list[LinkerEntity]): str
+ }
+
+ class SentenceChunkingStrategy <> extends chunking_strategies.ChunkingStrategy {
+ + chunk(document: ParsedDocument, doc_entity: DocumentAsEntity): list[LinkerEntity]
+ + dechunk(entities: list[LinkerEntity], links: list[LinkerEntity]): str
+ }
+
+ class NumberedItemsChunkingStrategy <> extends chunking_strategies.ChunkingStrategy {
+ + chunk(document: ParsedDocument, doc_entity: DocumentAsEntity): list[LinkerEntity]
+ + dechunk(entities: list[LinkerEntity], links: list[LinkerEntity]): str
+ }
+ }
+
+ note right of ChunkingStrategy
+ Базовая реализация dechunk сортирует чанки по chunk_index.
+ Стратегии могут переопределить, если им нужна
+ специфическая логика сборки
+ end note
+}
+
+package "core" {
+ class Destructurer <> {
+ + __init__(document: ParsedDocument, strategy_name: str)
+ + configure(strategy_name: str, **kwargs)
+ + destructure(): list[LinkerEntity]
+ }
+
+ class InjectionBuilder <> {
+ + __init__(entities: list[LinkerEntity], config: dict)
+ + register_strategy(doc_type: str, strategy: ChunkingStrategy)
+ + build(filtered_entities: list[LinkerEntity]): str
+ - _group_chunks_by_document(chunks, links): dict
+ }
+
+ note right of Destructurer
+ Основной класс библиотеки, используется для разбиения
+ документа на чанки и вспомогательные сущности. В
+ полученной конфигурации содержатся in_search сущности
+ и множество вспомогательных сущностей. Предполагается,
+ что первые будут отфильтрованы векторным или иным поиском,
+ а вторые можно будет использовать для обогащения и сборки
+ итоговой инъекции в промпт.
+ end note
+
+ note right of InjectionBuilder
+ Класс-единая точка входа для сборки итоговой инъекции
+ в промпт. Принимает в себя все сущности и конфигурацию
+ в конструкторе, а в методе build принимает отфильтрованные
+ сущности. Может частично делегировать сборку стратегиям для
+ специфических типов чанкинга.
+ end note
+
+}
+
+' Композиционные отношения
+core.Destructurer --> chunking_strategies.ChunkingStrategy
+core.InjectionBuilder --> chunking_strategies.ChunkingStrategy
+
+' Отношения между компонентами
+chunking_strategies.ChunkingStrategy ..> models
+
+' Дополнительные отношения
+core.InjectionBuilder ..> models.LinkerEntity
+core.Destructurer ..> models.LinkerEntity
+
+@enduml
\ No newline at end of file
diff --git a/lib/extractor/ntr_text_fragmentation/__init__.py b/lib/extractor/ntr_text_fragmentation/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..af29834fc7fb3c7c4b4218fbaeab8b38fc2a2e69
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/__init__.py
@@ -0,0 +1,19 @@
+"""
+Модуль извлечения и сборки документов.
+"""
+
+from .core.destructurer import Destructurer
+from .core.entity_repository import EntityRepository, InMemoryEntityRepository
+from .core.injection_builder import InjectionBuilder
+from .models import Chunk, DocumentAsEntity, LinkerEntity
+
+__all__ = [
+ "Destructurer",
+ "InjectionBuilder",
+ "EntityRepository",
+ "InMemoryEntityRepository",
+ "LinkerEntity",
+ "Chunk",
+ "DocumentAsEntity",
+ "integrations",
+]
diff --git a/lib/extractor/ntr_text_fragmentation/additors/__init__.py b/lib/extractor/ntr_text_fragmentation/additors/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae44f8c00be127f0339c321f4a399f5aed2239c1
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/additors/__init__.py
@@ -0,0 +1,10 @@
+"""
+Модуль для дополнительных обработчиков документа.
+
+Содержит обработчики, которые извлекают дополнительные сущности из документа,
+например, таблицы, изображения и т.д.
+"""
+
+from .tables_processor import TablesProcessor
+
+__all__ = ["TablesProcessor"]
diff --git a/lib/extractor/ntr_text_fragmentation/additors/tables/__init__.py b/lib/extractor/ntr_text_fragmentation/additors/tables/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f163392c046ac2890faf8800c7a80c24b0b66db5
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/additors/tables/__init__.py
@@ -0,0 +1,5 @@
+from .table_entity import TableEntity
+
+__all__ = [
+ 'TableEntity',
+]
diff --git a/lib/extractor/ntr_text_fragmentation/additors/tables/table_entity.py b/lib/extractor/ntr_text_fragmentation/additors/tables/table_entity.py
new file mode 100644
index 0000000000000000000000000000000000000000..e5792a83a31917e97798b012cadc26743e0509be
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/additors/tables/table_entity.py
@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from typing import Optional
+from uuid import UUID
+
+from ...models import LinkerEntity
+from ...models.linker_entity import register_entity
+
+
+@register_entity
+@dataclass
+class TableEntity(LinkerEntity):
+ """
+ Сущность таблицы из документа.
+
+ Расширяет основную сущность LinkerEntity, добавляя информацию о таблице.
+ """
+
+ table_index: Optional[int] = None
+
+ @classmethod
+ def deserialize(cls, entity: LinkerEntity) -> "TableEntity":
+ """
+ Десериализует сущность из базового LinkerEntity.
+
+ Args:
+ entity: Базовая сущность LinkerEntity
+
+ Returns:
+ Десериализованная сущность TableEntity
+ """
+ if entity.type != cls.__name__:
+ raise ValueError(f"Неверный тип сущности: {entity.type}, ожидался {cls.__name__}")
+
+ # Извлекаем дополнительные поля из метаданных
+ metadata = entity.metadata or {}
+ table_index = metadata.get("table_index")
+
+ return cls(
+ id=entity.id if isinstance(entity.id, UUID) else UUID(entity.id),
+ name=entity.name,
+ text=entity.text,
+ in_search_text=entity.in_search_text,
+ metadata=entity.metadata,
+ source_id=entity.source_id,
+ target_id=entity.target_id,
+ number_in_relation=entity.number_in_relation,
+ type=entity.type,
+ table_index=table_index,
+ )
+
+ def serialize(self) -> LinkerEntity:
+ """
+ Сериализует сущность в базовый LinkerEntity.
+
+ Returns:
+ Сериализованная сущность LinkerEntity
+ """
+ metadata = self.metadata or {}
+
+ # Добавляем дополнительные поля в метаданные
+ if self.table_index is not None:
+ metadata["table_index"] = self.table_index
+
+ return LinkerEntity(
+ id=self.id,
+ name=self.name,
+ text=self.text,
+ in_search_text=self.in_search_text,
+ metadata=metadata,
+ source_id=self.source_id,
+ target_id=self.target_id,
+ number_in_relation=self.number_in_relation,
+ type=self.__class__.__name__,
+ )
diff --git a/lib/extractor/ntr_text_fragmentation/additors/tables_processor.py b/lib/extractor/ntr_text_fragmentation/additors/tables_processor.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc40b87322292ce9e5a7d490dbb6f79cc4c5cebc
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/additors/tables_processor.py
@@ -0,0 +1,117 @@
+"""
+Процессор таблиц из документа.
+"""
+
+from uuid import uuid4
+
+from ntr_fileparser import ParsedDocument
+
+from ..models import LinkerEntity
+from .tables import TableEntity
+
+
+class TablesProcessor:
+ """
+ Процессор для извлечения таблиц из документа и создания связанных сущностей.
+ """
+
+ def __init__(self):
+ """Инициализация процессора таблиц."""
+ pass
+
+ def process(
+ self,
+ document: ParsedDocument,
+ doc_entity: LinkerEntity,
+ ) -> list[LinkerEntity]:
+ """
+ Извлекает таблицы из документа и создает для них сущности.
+
+ Args:
+ document: Документ для обработки
+ doc_entity: Сущность документа для связи с таблицами
+
+ Returns:
+ Список сущностей TableEntity и связей
+ """
+ if not document.tables:
+ return []
+
+ table_entities = []
+ links = []
+
+ rows = '\n\n'.join([table.to_string() for table in document.tables]).split(
+ '\n\n'
+ )
+
+ # Обрабатываем каждую таблицу
+ for idx, row in enumerate(rows):
+ # Создаем сущность таблицы
+ table_entity = self._create_table_entity(
+ table_text=row,
+ table_index=idx,
+ doc_name=doc_entity.name,
+ )
+
+ # Создаем связь между документом и таблицей
+ link = self._create_link(doc_entity, table_entity, idx)
+
+ table_entities.append(table_entity)
+ links.append(link)
+
+ # Возвращаем список таблиц и связей
+ return table_entities + links
+
+ def _create_table_entity(
+ self,
+ table_text: str,
+ table_index: int,
+ doc_name: str,
+ ) -> TableEntity:
+ """
+ Создает сущность таблицы.
+
+ Args:
+ table_text: Текст таблицы
+ table_index: Индекс таблицы в документе
+ doc_name: Имя документа
+
+ Returns:
+ Сущность TableEntity
+ """
+ entity_name = f"{doc_name}_table_{table_index}"
+
+ return TableEntity(
+ id=uuid4(),
+ name=entity_name,
+ text=table_text,
+ in_search_text=table_text,
+ metadata={},
+ type=TableEntity.__name__,
+ table_index=table_index,
+ )
+
+ def _create_link(
+ self, doc_entity: LinkerEntity, table_entity: TableEntity, index: int
+ ) -> LinkerEntity:
+ """
+ Создает связь между документом и таблицей.
+
+ Args:
+ doc_entity: Сущность документа
+ table_entity: Сущность таблицы
+ index: Индекс таблицы в документе
+
+ Returns:
+ Объект связи LinkerEntity
+ """
+ return LinkerEntity(
+ id=uuid4(),
+ name="document_to_table",
+ text="",
+ metadata={},
+ source_id=doc_entity.id,
+ target_id=table_entity.id,
+ number_in_relation=index,
+ type="Link",
+ )
diff --git a/lib/extractor/ntr_text_fragmentation/chunking/__init__.py b/lib/extractor/ntr_text_fragmentation/chunking/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d616841c4d30ce26f8be0214dcc79d7b815ef5c
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/chunking/__init__.py
@@ -0,0 +1,11 @@
+"""
+Модуль для определения стратегий чанкинга.
+"""
+
+from .chunking_strategy import ChunkingStrategy
+from .specific_strategies import FixedSizeChunkingStrategy
+
+__all__ = [
+ "ChunkingStrategy",
+ "FixedSizeChunkingStrategy",
+]
diff --git a/lib/extractor/ntr_text_fragmentation/chunking/chunking_strategy.py b/lib/extractor/ntr_text_fragmentation/chunking/chunking_strategy.py
new file mode 100644
index 0000000000000000000000000000000000000000..65f79be59d467c377c9465a6dccc28c5e9061292
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/chunking/chunking_strategy.py
@@ -0,0 +1,86 @@
+"""
+Базовый класс для всех стратегий чанкинга.
+"""
+
+from abc import ABC, abstractmethod
+
+from ntr_fileparser import ParsedDocument
+
+from ..models import Chunk, DocumentAsEntity, LinkerEntity
+
+
+class ChunkingStrategy(ABC):
+ """
+ Базовый абстрактный класс для всех стратегий чанкинга.
+ """
+
+ @abstractmethod
+ def chunk(self, document: ParsedDocument, doc_entity: DocumentAsEntity | None = None) -> list[LinkerEntity]:
+ """
+ Разбивает документ на чанки в соответствии со стратегией.
+
+ Args:
+ document: ParsedDocument для извлечения текста
+ doc_entity: Опциональная сущность документа для привязки чанков.
+ Если не указана, будет создана новая.
+
+ Returns:
+ list[LinkerEntity]: Список сущностей (документ, чанки, связи)
+ """
+ raise NotImplementedError("Стратегия чанкинга должна реализовать метод chunk")
+
+ def dechunk(self, chunks: list[LinkerEntity], repository: 'EntityRepository' = None) -> str:
+ """
+ Собирает документ из чанков и связей.
+
+ Базовая реализация сортирует чанки по chunk_index и объединяет их тексты,
+ сохраняя структуру параграфов и избегая дублирования текста.
+
+ Args:
+ chunks: Список отфильтрованных чанков в случайном порядке
+ repository: Репозиторий сущностей для получения дополнительной информации (может быть None)
+
+ Returns:
+ Восстановленный текст документа
+ """
+ import re
+
+ # Проверяем, есть ли чанки для сборки
+ if not chunks:
+ return ""
+
+ # Отбираем только чанки
+ valid_chunks = [c for c in chunks if isinstance(c, Chunk)]
+
+ # Сортируем чанки по chunk_index
+ sorted_chunks = sorted(valid_chunks, key=lambda c: c.chunk_index or 0)
+
+ # Собираем текст документа с учетом структуры параграфов
+ result_text = ""
+
+ for chunk in sorted_chunks:
+ # Получаем текст чанка (предпочитаем text, а не in_search_text для избежания дублирования)
+ chunk_text = chunk.text if hasattr(chunk, 'text') and chunk.text else ""
+
+ # Добавляем текст чанка с сохранением структуры параграфов
+ if result_text and result_text[-1] != "\n" and chunk_text and chunk_text[0] != "\n":
+ result_text += " "
+ result_text += chunk_text
+
+ # Пост-обработка результата
+ # Заменяем множественные переносы строк на одиночные
+ result_text = re.sub(r'\n+', '\n', result_text)
+
+ # Заменяем множественные пробелы на одиночные
+ result_text = re.sub(r' +', ' ', result_text)
+
+ # Убираем пробелы перед переносами строк
+ result_text = re.sub(r' +\n', '\n', result_text)
+
+ # Убираем пробелы после переносов строк
+ result_text = re.sub(r'\n +', '\n', result_text)
+
+ # Убираем лишние переносы строк в начале и конце текста
+ result_text = result_text.strip()
+
+ return result_text
\ No newline at end of file
diff --git a/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/__init__.py b/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..72bd3bf174484cee47cf1e77675050df527458ea
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/__init__.py
@@ -0,0 +1,11 @@
+"""
+Модуль содержащий конкретные стратегии для чанкинга текста.
+"""
+
+from .fixed_size import FixedSizeChunk
+from .fixed_size_chunking import FixedSizeChunkingStrategy
+
+__all__ = [
+ "FixedSizeChunk",
+ "FixedSizeChunkingStrategy",
+]
diff --git a/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size/__init__.py b/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f20ac6a4aa5d2841b3317ffdb53abf991fabf6e
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size/__init__.py
@@ -0,0 +1,9 @@
+"""
+Модуль реализующий стратегию чанкинга с фиксированным размером.
+"""
+
+from .fixed_size_chunk import FixedSizeChunk
+
+__all__ = [
+ "FixedSizeChunk",
+]
\ No newline at end of file
diff --git a/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size/fixed_size_chunk.py b/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size/fixed_size_chunk.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b00d808dd0e1e2cc29bc6dcfe6bcecffd8274b3
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size/fixed_size_chunk.py
@@ -0,0 +1,143 @@
+"""
+Класс для представления чанка фиксированного размера.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from ....models.chunk import Chunk
+from ....models.linker_entity import LinkerEntity, register_entity
+
+
+@register_entity
+@dataclass
+class FixedSizeChunk(Chunk):
+ """
+ Представляет чанк фиксированного размера.
+
+ Расширяет базовый класс Chunk дополнительными полями, связанными с токенами
+ и перекрытиями, а также добавляет методы для сборки документа с учетом
+ границ предложений.
+ """
+
+ token_count: int = 0
+
+ # Информация о границах предложений и нахлестах
+ left_sentence_part: str = "" # Часть предложения слева от text
+ right_sentence_part: str = "" # Часть предложения справа от text
+ overlap_left: str = "" # Нахлест слева (без учета границ предложений)
+ overlap_right: str = "" # Нахлест справа (без учета границ предложений)
+
+ # Метаданные для дополнительной информации
+ metadata: dict[str, Any] = field(default_factory=dict)
+
+ def __str__(self) -> str:
+ """
+ Строковое представление чанка.
+
+ Returns:
+ Строка с информацией о чанке.
+ """
+ return (
+ f"FixedSizeChunk(id={self.id}, chunk_index={self.chunk_index}, "
+ f"tokens={self.token_count}, "
+ f"text='{self.text[:30]}{'...' if len(self.text) > 30 else ''}'"
+ f")"
+ )
+
+ def get_adjacent_chunks_indices(self, max_distance: int = 1) -> list[int]:
+ """
+ Возвращает индексы соседних чанков в пределах указанного расстояния.
+
+ Args:
+ max_distance: Максимальное расстояние от текущего чанка
+
+ Returns:
+ Список индексов соседних чанков
+ """
+ indices = []
+ for i in range(1, max_distance + 1):
+ # Добавляем предыдущие чанки
+ if self.chunk_index - i >= 0:
+ indices.append(self.chunk_index - i)
+ # Добавляем следующие чанки
+ indices.append(self.chunk_index + i)
+
+ return sorted(indices)
+
+ @classmethod
+ def deserialize(cls, entity: LinkerEntity) -> 'FixedSizeChunk':
+ """
+ Десериализует FixedSizeChunk из объекта LinkerEntity.
+
+ Args:
+ entity: Объект LinkerEntity для преобразования в FixedSizeChunk
+
+ Returns:
+ Десериализованный объект FixedSizeChunk
+ """
+ metadata = entity.metadata or {}
+
+ # Извлекаем параметры из метаданных
+ # Сначала проверяем в метаданных под ключом _chunk_index
+ chunk_index = metadata.get('_chunk_index')
+ if chunk_index is None:
+ # Затем пробуем получить как атрибут объекта
+ chunk_index = getattr(entity, 'chunk_index', None)
+ if chunk_index is None:
+ # Если и там нет, пробуем обычный поиск по метаданным
+ chunk_index = metadata.get('chunk_index')
+
+ # Преобразуем к int, если значение найдено
+ if chunk_index is not None:
+ try:
+ chunk_index = int(chunk_index)
+ except (ValueError, TypeError):
+ chunk_index = None
+
+ start_token = metadata.get('start_token', 0)
+ end_token = metadata.get('end_token', 0)
+ token_count = metadata.get(
+ '_token_count', metadata.get('token_count', end_token - start_token + 1)
+ )
+
+ # Извлекаем параметры для границ предложений и нахлестов
+ # Сначала ищем в метаданных с префиксом _
+ left_sentence_part = metadata.get('_left_sentence_part')
+ if left_sentence_part is None:
+ # Затем пробуем получить как атрибут объекта
+ left_sentence_part = getattr(entity, 'left_sentence_part', '')
+
+ right_sentence_part = metadata.get('_right_sentence_part')
+ if right_sentence_part is None:
+ right_sentence_part = getattr(entity, 'right_sentence_part', '')
+
+ overlap_left = metadata.get('_overlap_left')
+ if overlap_left is None:
+ overlap_left = getattr(entity, 'overlap_left', '')
+
+ overlap_right = metadata.get('_overlap_right')
+ if overlap_right is None:
+ overlap_right = getattr(entity, 'overlap_right', '')
+
+ # Создаем чистые метаданные без служебных полей
+ clean_metadata = {k: v for k, v in metadata.items() if not k.startswith('_')}
+
+ # Создаем и возвращаем новый экземпляр FixedSizeChunk
+ return cls(
+ id=entity.id,
+ name=entity.name,
+ text=entity.text,
+ in_search_text=entity.in_search_text,
+ metadata=clean_metadata,
+ source_id=entity.source_id,
+ target_id=entity.target_id,
+ number_in_relation=entity.number_in_relation,
+ chunk_index=chunk_index,
+ token_count=token_count,
+ left_sentence_part=left_sentence_part,
+ right_sentence_part=right_sentence_part,
+ overlap_left=overlap_left,
+ overlap_right=overlap_right,
+ type="FixedSizeChunk",
+ )
diff --git a/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size_chunking.py b/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size_chunking.py
new file mode 100644
index 0000000000000000000000000000000000000000..a3ccb9542b6103f6df605674a6dda52c901897f9
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/chunking/specific_strategies/fixed_size_chunking.py
@@ -0,0 +1,568 @@
+"""
+Стратегия чанкинга фиксированного размера.
+"""
+
+import re
+from typing import NamedTuple, TypeVar
+from uuid import uuid4
+
+from ntr_fileparser import ParsedDocument, ParsedTextBlock
+
+from ...chunking.chunking_strategy import ChunkingStrategy
+from ...models import DocumentAsEntity, LinkerEntity
+from .fixed_size.fixed_size_chunk import FixedSizeChunk
+
+T = TypeVar('T')
+
+
+class _FixedSizeChunkingStrategyParams(NamedTuple):
+ words_per_chunk: int = 50
+ overlap_words: int = 25
+ respect_sentence_boundaries: bool = True
+
+
+class FixedSizeChunkingStrategy(ChunkingStrategy):
+ """
+ Стратегия чанкинга, разбивающая текст на чанки фиксированного размера.
+
+ Преимущества:
+ - Простое и предсказуемое разбиение
+ - Равные по размеру чанки
+
+ Недостатки:
+ - Может разрезать предложения и абзацы в середине (компенсируется сборкой - как для модели поиска, так и для LLM)
+ - Не учитывает смысловую структуру текста
+
+ Особенности реализации:
+ - В поле `text` чанков хранится текст без нахлеста (для удобства сборки)
+ - В поле `in_search_text` хранится текст с нахлестом (для улучшения векторизации)
+ """
+
+ name = "fixed_size"
+ description = (
+ "Стратегия чанкинга, разбивающая текст на чанки фиксированного размера."
+ )
+
+ def __init__(
+ self,
+ words_per_chunk: int = 50,
+ overlap_words: int = 25,
+ respect_sentence_boundaries: bool = True,
+ ):
+ """
+ Инициализация стратегии чанкинга с фиксированным размером.
+
+ Args:
+ words_per_chunk: Количество слов в чанке
+ overlap_words: Количество слов перекрытия между чанками
+ respect_sentence_boundaries: Флаг учета границ предложений
+ """
+
+ self.params = _FixedSizeChunkingStrategyParams(
+ words_per_chunk=words_per_chunk,
+ overlap_words=overlap_words,
+ respect_sentence_boundaries=respect_sentence_boundaries,
+ )
+
+ def chunk(
+ self,
+ document: ParsedDocument | str,
+ doc_entity: DocumentAsEntity | None = None,
+ ) -> list[LinkerEntity]:
+ """
+ Разбивает документ на чанки фиксированного размера.
+
+ Args:
+ document: Документ для разбиения (ParsedDocument или текст)
+ doc_entity: Сущность документа (опционально)
+
+ Returns:
+ Список LinkerEntity - чанки, связи и прочие сущности
+ """
+ doc = self._prepare_document(document)
+ words = self._extract_words(doc)
+
+ # Если документ пустой, возвращаем пустой список
+ if not words:
+ return []
+
+ doc_entity = self._ensure_document_entity(doc, doc_entity)
+ doc_name = doc_entity.name
+
+ chunks = []
+ links = []
+
+ step = self._calculate_step()
+ total_words = len(words)
+
+ # Начинаем с первого слова и идем шагами (не полным размером чанка)
+ for i in range(0, total_words, step):
+ # Создаем обычный чанк
+ chunk_text = self._prepare_chunk_text(words, i, step)
+ in_search_text = self._prepare_chunk_text(
+ words, i, self.params.words_per_chunk
+ )
+
+ chunk = self._create_chunk(
+ chunk_text,
+ in_search_text,
+ i,
+ i + self.params.words_per_chunk,
+ len(chunks),
+ words,
+ total_words,
+ doc_name,
+ )
+
+ chunks.append(chunk)
+ links.append(self._create_link(doc_entity, chunk))
+
+ # Возвращаем все сущности
+ return [doc_entity] + chunks + links
+
+ def _find_nearest_sentence_boundary(
+ self, text: str, position: int
+ ) -> tuple[int, str, str]:
+ """
+ Находит ближайшую границу предложения к указанной позиции.
+
+ Args:
+ text: Полный текст для поиска границ
+ position: Позиция, для которой ищем ближайшую границу
+
+ Returns:
+ tuple из (позиция границы, левая часть текста, правая часть текста)
+ """
+ # Регулярное выражение для поиска конца предложения
+ sentence_end_pattern = r'[.!?](?:\s|$)'
+
+ # Ищем все совпадения в тексте
+ matches = list(re.finditer(sentence_end_pattern, text))
+
+ if not matches:
+ # Если совпадений нет, возвращаем исходную позицию
+ return position, text[:position], text[position:]
+
+ # Находим ближайшую границу предложения
+ nearest_pos = position
+ min_distance = float('inf')
+
+ for match in matches:
+ end_pos = match.end()
+ distance = abs(end_pos - position)
+
+ if distance < min_distance:
+ min_distance = distance
+ nearest_pos = end_pos
+
+ # Возвращаем позицию и соответствующие части текста
+ return nearest_pos, text[:nearest_pos], text[nearest_pos:]
+
+ def _find_sentence_boundary(self, text: str, is_left_boundary: bool) -> str:
+ """
+ Находит часть текста на границе предложения.
+
+ Args:
+ text: Текст для обработки
+ is_left_boundary: True для левой границы, False для правой
+
+ Returns:
+ Часть предложения на границе
+ """
+ # Регулярное выражение для поиска конца предложения
+ sentence_end_pattern = r'[.!?](?:\s|$)'
+ matches = list(re.finditer(sentence_end_pattern, text))
+
+ if not matches:
+ return text
+
+ if is_left_boundary:
+ # Для левой границы берем часть после последней границы предложения
+ last_match = matches[-1]
+ return text[last_match.end() :].strip()
+ else:
+ # Для правой границы берем часть до первой границы предложения
+ first_match = matches[0]
+ return text[: first_match.end()].strip()
+
+ def dechunk(
+ self,
+ filtered_chunks: list[LinkerEntity],
+ repository: 'EntityRepository' = None, # type: ignore
+ ) -> str:
+ """
+ Собирает документ из чанков и связей.
+
+ Args:
+ filtered_chunks: Список отфильтрованных чанков
+ repository: Репозиторий сущностей для получения дополнительной информации (может быть None)
+
+ Returns:
+ Восстановленный текст документа
+ """
+ if not filtered_chunks:
+ return ""
+
+ # Проверяем тип и десериализуем FixedSizeChunk
+ chunks = []
+ for chunk in filtered_chunks:
+ if chunk.type == "FixedSizeChunk":
+ chunks.append(FixedSizeChunk.deserialize(chunk))
+ else:
+ chunks.append(chunk)
+
+ # Сортируем чанки по индексу
+ sorted_chunks = sorted(chunks, key=lambda c: c.chunk_index or 0)
+
+ # Инициализируем результирующий текст
+ result_text = ""
+
+ # Группируем последовательные чанки
+ current_group = []
+ groups = []
+
+ for i, chunk in enumerate(sorted_chunks):
+ current_index = chunk.chunk_index or 0
+
+ # Если первый чанк или продолжение последовательности
+ if i == 0 or current_index == (sorted_chunks[i - 1].chunk_index or 0) + 1:
+ current_group.append(chunk)
+ else:
+ # Закрываем текущую группу и начинаем новую
+ if current_group:
+ groups.append(current_group)
+ current_group = [chunk]
+
+ # Добавляем последнюю группу
+ if current_group:
+ groups.append(current_group)
+
+ # Обрабатываем каждую группу
+ for group_index, group in enumerate(groups):
+ # Добавляем многоточие между непоследовательными группами
+ if group_index > 0:
+ result_text += "\n\n...\n\n"
+
+ # Обрабатываем группу соседних чанков
+ group_text = ""
+
+ # Добавляем левую недостающую часть к первому чанку группы
+ first_chunk = group[0]
+
+ # Добавляем левую часть предложения к первому чанку группы
+ if (
+ hasattr(first_chunk, 'left_sentence_part')
+ and first_chunk.left_sentence_part
+ ):
+ group_text += first_chunk.left_sentence_part
+
+ # Добавляем текст всех чанков группы
+ for i, chunk in enumerate(group):
+ current_text = chunk.text.strip() if hasattr(chunk, 'text') else ""
+ if not current_text:
+ continue
+
+ # Проверяем, нужно ли добавить пробел между предыдущим текстом и текущим чанком
+ if group_text:
+ # Если текущий чанк начинается с новой строки, не добавляем пробел
+ if current_text.startswith("\n"):
+ pass
+ # Если предыдущий текст заканчивается переносом строки, также не добавляем пробел
+ elif group_text.endswith("\n"):
+ pass
+ # Если предыдущий текст заканчивается знаком препинания без пробела, добавляем пробел
+ elif group_text.rstrip()[-1] not in [
+ "\n",
+ " ",
+ ".",
+ ",",
+ "!",
+ "?",
+ ":",
+ ";",
+ "-",
+ "–",
+ "—",
+ ]:
+ group_text += " "
+
+ # Добавляем текст чанка
+ group_text += current_text
+
+ # Добавляем правую недостающую часть к последнему чанку группы
+ last_chunk = group[-1]
+
+ # Добавляем правую часть предложения к последнему чанку группы
+ if (
+ hasattr(last_chunk, 'right_sentence_part')
+ and last_chunk.right_sentence_part
+ ):
+ right_part = last_chunk.right_sentence_part.strip()
+ if right_part:
+ # Проверяем нужен ли пробел перед правой частью
+ if (
+ group_text
+ and group_text[-1] not in ["\n", " "]
+ and right_part[0]
+ not in ["\n", " ", ".", ",", "!", "?", ":", ";", "-", "–", "—"]
+ ):
+ group_text += " "
+ group_text += right_part
+
+ # Добавляем текст группы к результату
+ if (
+ result_text
+ and result_text[-1] not in ["\n", " "]
+ and group_text
+ and group_text[0] not in ["\n", " "]
+ ):
+ result_text += " "
+ result_text += group_text
+
+ # Постобработка текста: удаляем лишние пробелы и символы переноса строк
+
+ # Заменяем множественные переносы строк на двойные (для разделения абзацев)
+ result_text = re.sub(r'\n{3,}', '\n\n', result_text)
+
+ # Заменяем множественные пробелы на одиночные
+ result_text = re.sub(r' +', ' ', result_text)
+
+ # Убираем пробелы перед знаками препинания
+ result_text = re.sub(r' ([.,!?:;)])', r'\1', result_text)
+
+ # Убираем пробелы перед переносами строк и после переносов строк
+ result_text = re.sub(r' +\n', '\n', result_text)
+ result_text = re.sub(r'\n +', '\n', result_text)
+
+ # Убираем лишние переносы строк и пробелы в начале и конце текста
+ result_text = result_text.strip()
+
+ return result_text
+
+ def _get_sorted_chunks(
+ self, chunks: list[LinkerEntity], links: list[LinkerEntity]
+ ) -> list[LinkerEntity]:
+ """
+ Получает отсортированные чанки на основе связей или поля chunk_index.
+
+ Args:
+ chunks: Список чанков для сортировки
+ links: Список связей для определения порядка
+
+ Returns:
+ Отсортированные чанки
+ """
+ # Сортируем чанки по порядку в связях
+ if links:
+ # Получаем словарь для быстрого доступа к чанкам по ID
+ chunk_dict = {c.id: c for c in chunks}
+
+ # Сортируем по порядку в связях
+ sorted_chunks = []
+ for link in sorted(links, key=lambda l: l.number_in_relation or 0):
+ if link.target_id in chunk_dict:
+ sorted_chunks.append(chunk_dict[link.target_id])
+
+ return sorted_chunks
+
+ # Если нет связей, сортируем по chunk_index
+ return sorted(chunks, key=lambda c: c.chunk_index or 0)
+
+ def _prepare_document(self, document: ParsedDocument | str) -> ParsedDocument:
+ """
+ Обрабатывает входные данные и возвращает ParsedDocument.
+
+ Args:
+ document: Документ (ParsedDocument или текст)
+
+ Returns:
+ Обработанный документ типа ParsedDocument
+ """
+ if isinstance(document, ParsedDocument):
+ return document
+ elif isinstance(document, str):
+ # Простая обработка текстового документа
+ return ParsedDocument(
+ paragraphs=[
+ ParsedTextBlock(text=paragraph)
+ for paragraph in document.split('\n')
+ ]
+ )
+
+ def _extract_words(self, doc: ParsedDocument) -> list[str]:
+ """
+ Извлекает все слова из документа.
+
+ Args:
+ doc: Документ для извлечения слов
+
+ Returns:
+ Список слов документа
+ """
+ words = []
+ for paragraph in doc.paragraphs:
+ # Добавляем слова из параграфа
+ paragraph_words = paragraph.text.split()
+ words.extend(paragraph_words)
+ # Добавляем маркер конца параграфа как отдельный элемент
+ words.append("\n")
+ return words
+
+ def _ensure_document_entity(
+ self,
+ doc: ParsedDocument,
+ doc_entity: LinkerEntity | None,
+ ) -> LinkerEntity:
+ """
+ Создает сущность документа, если не предоставлена.
+
+ Args:
+ doc: Документ
+ doc_entity: Сущность документа (может быть None)
+
+ Returns:
+ Сущность документа
+ """
+ if doc_entity is None:
+ return LinkerEntity(
+ id=uuid4(),
+ name=doc.name,
+ text=doc.name,
+ metadata={"type": doc.type},
+ type="Document",
+ )
+ return doc_entity
+
+ def _calculate_step(self) -> int:
+ """
+ Вычисляет шаг для создания чанков.
+
+ Returns:
+ Размер шага между началами чанков
+ """
+ return self.params.words_per_chunk - self.params.overlap_words
+
+ def _prepare_chunk_text(
+ self,
+ words: list[str],
+ start_idx: int,
+ length: int,
+ ) -> str:
+ """
+ Подготавливает текст чанка и текст для поиска.
+
+ Args:
+ words: Список слов документа
+ start_idx: Индекс начала чанка
+ end_idx: Длина текста в словах
+
+ Returns:
+ Итоговый текст
+ """
+ # Извлекаем текст чанка без нахлеста с сохранением структуры параграфов
+ end_idx = min(start_idx + length, len(words))
+ chunk_words = words[start_idx:end_idx]
+ chunk_text = ""
+
+ for word in chunk_words:
+ if word == "\n":
+ # Если это маркер конца параграфа, добавляем перенос строки
+ chunk_text += "\n"
+ else:
+ # Иначе добавляем слово с пробелом
+ if chunk_text and chunk_text[-1] != "\n":
+ chunk_text += " "
+ chunk_text += word
+
+ return chunk_text
+
+ def _create_chunk(
+ self,
+ chunk_text: str,
+ in_search_text: str,
+ start_idx: int,
+ end_idx: int,
+ chunk_index: int,
+ words: list[str],
+ total_words: int,
+ doc_name: str,
+ ) -> FixedSizeChunk:
+ """
+ Создает чанк фиксированного размера.
+
+ Args:
+ chunk_text: Текст чанка без нахлеста
+ in_search_text: Текст чанка с нахлестом
+ start_idx: Индекс первого слова в чанке
+ end_idx: Индекс последнего слова в чанке
+ chunk_index: Индекс чанка в документе
+ words: Список всех слов документа
+ total_words: Общее количество слов в документе
+ doc_name: Имя документа
+
+ Returns:
+ FixedSizeChunk: Созданный чанк
+ """
+ # Определяем нахлесты без учета границ предложений
+ overlap_left = " ".join(
+ words[max(0, start_idx - self.params.overlap_words) : start_idx]
+ )
+ overlap_right = " ".join(
+ words[end_idx : min(total_words, end_idx + self.params.overlap_words)]
+ )
+
+ # Определяем границы предложений
+ left_sentence_part = ""
+ right_sentence_part = ""
+
+ if self.params.respect_sentence_boundaries:
+ # Находим ближайшую границу предложения слева
+ left_text = " ".join(
+ words[max(0, start_idx - self.params.overlap_words) : start_idx]
+ )
+ left_sentence_part = self._find_sentence_boundary(left_text, True)
+
+ # Находим ближайшую границу предложения справа
+ right_text = " ".join(
+ words[end_idx : min(total_words, end_idx + self.params.overlap_words)]
+ )
+ right_sentence_part = self._find_sentence_boundary(right_text, False)
+
+ # Создаем чанк с учетом границ предложений
+ return FixedSizeChunk(
+ id=uuid4(),
+ name=f"{doc_name}_chunk_{chunk_index}",
+ text=chunk_text,
+ chunk_index=chunk_index,
+ in_search_text=in_search_text,
+ token_count=end_idx - start_idx + 1,
+ left_sentence_part=left_sentence_part,
+ right_sentence_part=right_sentence_part,
+ overlap_left=overlap_left,
+ overlap_right=overlap_right,
+ metadata={},
+ type=FixedSizeChunk.__name__,
+ )
+
+ def _create_link(
+ self, doc_entity: LinkerEntity, chunk: LinkerEntity
+ ) -> LinkerEntity:
+ """
+ Создает связь между документом и чанком.
+
+ Args:
+ doc_entity: Сущность документа
+ chunk: Сущность чанка
+
+ Returns:
+ Объект связи
+ """
+ return LinkerEntity(
+ id=uuid4(),
+ name="document_to_chunk",
+ text="",
+ metadata={},
+ source_id=doc_entity.id,
+ target_id=chunk.id,
+ type="Link",
+ )
diff --git a/lib/extractor/ntr_text_fragmentation/core/__init__.py b/lib/extractor/ntr_text_fragmentation/core/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8520f88d256d37ddeec87769b3269be23316dbb6
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/core/__init__.py
@@ -0,0 +1,9 @@
+"""
+Основные классы для разбиения и сборки документов.
+"""
+
+from .destructurer import Destructurer
+from .entity_repository import EntityRepository, InMemoryEntityRepository
+from .injection_builder import InjectionBuilder
+
+__all__ = ["Destructurer", "InjectionBuilder", "EntityRepository", "InMemoryEntityRepository"]
diff --git a/lib/extractor/ntr_text_fragmentation/core/destructurer.py b/lib/extractor/ntr_text_fragmentation/core/destructurer.py
new file mode 100644
index 0000000000000000000000000000000000000000..48638f7262a342bd340386a068d70d0b0fa36682
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/core/destructurer.py
@@ -0,0 +1,143 @@
+"""
+Модуль для деструктуризации документа.
+"""
+
+from uuid import uuid4
+
+# Внешние импорты
+from ntr_fileparser import ParsedDocument
+
+# Импорты из этой же библиотеки
+from ..additors.tables_processor import TablesProcessor
+from ..chunking.chunking_strategy import ChunkingStrategy
+from ..chunking.specific_strategies.fixed_size_chunking import \
+ FixedSizeChunkingStrategy
+from ..models import DocumentAsEntity, LinkerEntity
+
+
+class Destructurer:
+ """
+ Класс для подготовки документа для загрузки в базу данных.
+ Разбивает документ на чанки, создает связи между ними и
+ извлекает вспомогательные сущности.
+ """
+
+ # Доступные стратегии чанкинга
+ STRATEGIES: dict[str, type[ChunkingStrategy]] = {
+ "fixed_size": FixedSizeChunkingStrategy,
+ }
+
+ def __init__(
+ self,
+ document: ParsedDocument,
+ strategy_name: str = "fixed_size",
+ process_tables: bool = True,
+ **kwargs,
+ ):
+ """
+ Инициализация деструктуризатора.
+
+ Args:
+ document: Документ для обработки
+ strategy_name: Имя стратегии
+ process_tables: Флаг обработки таблиц
+ **kwargs: Параметры для стратегии
+ """
+ self.document = document
+ self.strategy: ChunkingStrategy | None = None
+ self.process_tables = process_tables
+ # Инициализируем процессор таблиц, если нужно
+ self.tables_processor = TablesProcessor() if process_tables else None
+ # Кеш для хранения созданных стратегий
+ self._strategy_cache: dict[str, ChunkingStrategy] = {}
+
+ # Конфигурируем стратегию
+ self.configure(strategy_name, **kwargs)
+
+ def configure(self, strategy_name: str = "fixed_size", **kwargs) -> None:
+ """
+ Установка стратегии чанкинга.
+
+ Args:
+ strategy_name: Имя стратегии
+ **kwargs: Параметры для стратегии
+
+ Raises:
+ ValueError: Если указана неизвестная стратегия
+ """
+ # Получаем класс стратегии из словаря доступных стратегий
+ if strategy_name not in self.STRATEGIES:
+ raise ValueError(f"Неизвестная стратегия: {strategy_name}")
+
+ # Создаем ключ кеша на основе имени стратегии и параметров
+ cache_key = f"{strategy_name}_{hash(frozenset(kwargs.items()))}"
+
+ # Проверяем, есть ли стратегия в кеше
+ if cache_key in self._strategy_cache:
+ self.strategy = self._strategy_cache[cache_key]
+ return
+
+ # Создаем экземпляр стратегии с переданными параметрами
+ strategy_class = self.STRATEGIES[strategy_name]
+ self.strategy = strategy_class(**kwargs)
+
+ # Сохраняем стратегию в кеше
+ self._strategy_cache[cache_key] = self.strategy
+
+ def destructure(self) -> list[LinkerEntity]:
+ """
+ Основной метод деструктуризации.
+ Разбивает документ на чанки и создает связи.
+
+ Returns:
+ list[LinkerEntity]: список сущностей, включая связи
+
+ Raises:
+ RuntimeError: Если стратегия не была сконфигурирована
+ """
+ # Проверяем, что стратегия сконфигурирована
+ if self.strategy is None:
+ raise RuntimeError("Стратегия не была сконфигурирована")
+
+ # Создаем сущность документа с метаданными
+ doc_entity = self._create_document_entity()
+
+ # Применяем стратегию чанкинга
+ entities = self.strategy.chunk(self.document, doc_entity)
+
+ # Обрабатываем таблицы, если это включено
+ if self.process_tables and self.tables_processor and self.document.tables:
+ table_entities = self.tables_processor.process(self.document, doc_entity)
+ entities.extend(table_entities)
+
+ # Сериализуем все сущности в простейшую форму LinkerEntity
+ serialized_entities = [entity.serialize() for entity in entities]
+
+ return serialized_entities
+
+ def _create_document_entity(self) -> DocumentAsEntity:
+ """
+ Создает сущность документа с метаданными.
+
+ Returns:
+ DocumentAsEntity: сущность документа
+ """
+ # Получаем имя документа или используем значение по умолчанию
+ doc_name = self.document.name or "Document"
+
+ # Создаем метаданные, включая информацию о стратегии чанкинга
+ metadata = {
+ "type": self.document.type,
+ "chunking_strategy": (
+ self.strategy.__class__.__name__ if self.strategy else "unknown"
+ ),
+ }
+
+ # Создаем сущность документа
+ return DocumentAsEntity(
+ id=uuid4(),
+ name=doc_name,
+ text="",
+ metadata=metadata,
+ type="Document",
+ )
diff --git a/lib/extractor/ntr_text_fragmentation/core/entity_repository.py b/lib/extractor/ntr_text_fragmentation/core/entity_repository.py
new file mode 100644
index 0000000000000000000000000000000000000000..96837733cb03f45bf75919f86a73b59cdbb6f474
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/core/entity_repository.py
@@ -0,0 +1,258 @@
+"""
+Интерфейс репозитория сущностей.
+"""
+
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import Iterable
+from uuid import UUID
+
+from ..models import Chunk, LinkerEntity
+from ..models.document import DocumentAsEntity
+
+
+class EntityRepository(ABC):
+ """
+ Абстрактный интерфейс для доступа к хранилищу сущностей.
+ Позволяет InjectionBuilder получать нужные сущности независимо от их хранилища.
+
+ Этот интерфейс определяет только методы для получения сущностей.
+ Логика сохранения и изменения сущностей остается за пределами этого интерфейса
+ и должна быть реализована в конкретных классах, расширяющих данный интерфейс.
+ """
+
+ @abstractmethod
+ def get_entities_by_ids(self, entity_ids: Iterable[UUID]) -> list[LinkerEntity]:
+ """
+ Получить сущности по списку идентификаторов.
+
+ Args:
+ entity_ids: Список идентификаторов сущностей
+
+ Returns:
+ Список сущностей, соответствующих указанным идентификаторам
+ """
+ pass
+
+ @abstractmethod
+ def get_document_for_chunks(self, chunk_ids: Iterable[UUID]) -> list[LinkerEntity]:
+ """
+ Получить документы, которым принадлежат указанные чанки.
+
+ Args:
+ chunk_ids: Список идентификаторов чанков
+
+ Returns:
+ Список документов, которым принадлежат указанные чанки
+ """
+ pass
+
+ @abstractmethod
+ def get_neighboring_chunks(self,
+ chunk_ids: Iterable[UUID],
+ max_distance: int = 1) -> list[LinkerEntity]:
+ """
+ Получить соседние чанки для указанных чанков.
+
+ Args:
+ chunk_ids: Список идентификаторов чанков
+ max_distance: Максимальное расстояние до соседа
+
+ Returns:
+ Список соседних чанков
+ """
+ pass
+
+ @abstractmethod
+ def get_related_entities(self,
+ entity_ids: Iterable[UUID],
+ relation_name: str | None = None,
+ as_source: bool = False,
+ as_target: bool = False) -> list[LinkerEntity]:
+ """
+ Получить сущности, связанные с указанными сущностями.
+
+ Args:
+ entity_ids: Список идентификаторов сущностей
+ relation_name: Опциональное имя отношения для фильтрации
+ as_source: Если True, ищем связи, где указанные entity_ids являются
+ источниками (source_id)
+ as_target: Если True, ищем связи, где указанные entity_ids являются
+ целевыми (target_id)
+
+ Returns:
+ Список связанных сущностей и связей
+ """
+ pass
+
+
+class InMemoryEntityRepository(EntityRepository):
+ """
+ Реализация EntityRepository, хранящая все сущности в памяти.
+ Обеспечивает обратную совместимость и используется для тестирования.
+ """
+
+ def __init__(self, entities: list[LinkerEntity] | None = None):
+ """
+ Инициализация репозитория с начальным списком сущностей.
+
+ Args:
+ entities: Начальный список сущностей
+ """
+ self.entities = entities or []
+ self._build_indices()
+
+ def _build_indices(self) -> None:
+ """
+ Строит индексы для быстрого доступа к сущностям.
+ """
+ self.entities_by_id = {e.id: e for e in self.entities}
+ self.chunks = [e for e in self.entities if isinstance(e, Chunk)]
+ self.docs = [e for e in self.entities if isinstance(e, DocumentAsEntity)]
+
+ # Индексы для быстрого поиска связей
+ self.doc_to_chunks = defaultdict(list)
+ self.chunk_to_doc = {}
+ self.entity_relations = defaultdict(list)
+ self.entity_targets = defaultdict(list)
+
+ # Заполняем индексы
+ for e in self.entities:
+ if e.is_link():
+ self.entity_relations[e.source_id].append(e)
+ self.entity_targets[e.target_id].append(e)
+ if e.name == "document_to_chunk":
+ self.doc_to_chunks[e.source_id].append(e.target_id)
+ self.chunk_to_doc[e.target_id] = e.source_id
+ if e.name == "document_to_table":
+ self.entity_relations
+ self.entity_targets[e.source_id].append(e.target_id)
+
+ # Этот метод не является частью интерфейса EntityRepository,
+ # но он полезен для тестирования и реализации обратной совместимости
+ def add_entities(self, entities: list[LinkerEntity]) -> None:
+ """
+ Добавляет сущности в репозиторий.
+
+ Примечание: Этот метод не является частью интерфейса EntityRepository.
+ Он добавлен для удобства тестирования и обратной совместимости.
+
+ Args:
+ entities: Список сущностей для добавления
+ """
+ self.entities.extend(entities)
+ self._build_indices()
+
+ def get_entities_by_ids(self, entity_ids: Iterable[UUID]) -> list[LinkerEntity]:
+ result = [self.entities_by_id.get(eid) for eid in entity_ids if eid in self.entities_by_id]
+ return result
+
+ def get_document_for_chunks(self, chunk_ids: Iterable[UUID]) -> list[LinkerEntity]:
+ result = []
+ for chunk_id in chunk_ids:
+ doc_id = self.chunk_to_doc.get(chunk_id)
+ if doc_id and doc_id in self.entities_by_id:
+ doc = self.entities_by_id[doc_id]
+ if doc not in result:
+ result.append(doc)
+ return result
+
+ def get_neighboring_chunks(self,
+ chunk_ids: Iterable[UUID],
+ max_distance: int = 1) -> list[LinkerEntity]:
+ result = []
+ chunk_indices = {}
+
+ # Сначала собираем индексы всех указанных чанков
+ for chunk_id in chunk_ids:
+ if chunk_id in self.entities_by_id:
+ chunk = self.entities_by_id[chunk_id]
+ if hasattr(chunk, 'chunk_index') and chunk.chunk_index is not None:
+ chunk_indices[chunk_id] = chunk.chunk_index
+
+ # Если нет чанков с индексами, возвращаем пустой список
+ if not chunk_indices:
+ return []
+
+ # Затем для каждого документа находим соседние чанки
+ for doc_id, doc_chunk_ids in self.doc_to_chunks.items():
+ # Проверяем, принадлежит ли хоть один из чанков этому документу
+ has_chunks = any(chunk_id in doc_chunk_ids for chunk_id in chunk_ids)
+ if not has_chunks:
+ continue
+
+ # Для каждого чанка в документе проверяем, является ли он соседом
+ for doc_chunk_id in doc_chunk_ids:
+ if doc_chunk_id in self.entities_by_id:
+ chunk = self.entities_by_id[doc_chunk_id]
+
+ # Если у чанка нет индекса, пропускаем его
+ if not hasattr(chunk, 'chunk_index') or chunk.chunk_index is None:
+ continue
+
+ # Проверяем, является ли чанк соседом какого-либо из исходных чанков
+ for orig_chunk_id, orig_index in chunk_indices.items():
+ if abs(chunk.chunk_index - orig_index) <= max_distance and doc_chunk_id not in chunk_ids:
+ result.append(chunk)
+ break
+
+ return result
+
+ def get_related_entities(self,
+ entity_ids: Iterable[UUID],
+ relation_name: str | None = None,
+ as_source: bool = False,
+ as_target: bool = False) -> list[LinkerEntity]:
+ """
+ Получить сущности, связанные с указанными сущностями.
+
+ Args:
+ entity_ids: Список идентификаторов сущностей
+ relation_name: Опциональное имя отношения для фильтрации
+ as_source: Если True, ищем связи, где указанные entity_ids являются источниками
+ as_target: Если True, ищем связи, где указанные entity_ids являются целями
+
+ Returns:
+ Список связанных сущностей и связей
+ """
+ result = []
+
+ # Если не указано ни as_source, ни as_target, по умолчанию ищем связи,
+ # где указанные entity_ids являются источниками
+ if not as_source and not as_target:
+ as_source = True
+
+ for entity_id in entity_ids:
+ if as_source:
+ # Ищем связи, где сущность является источником
+ relations = self.entity_relations.get(entity_id, [])
+
+ for link in relations:
+ if relation_name is None or link.name == relation_name:
+ # Добавляем саму связь
+ if link not in result:
+ result.append(link)
+
+ # Добавляем целевую сущность
+ if link.target_id in self.entities_by_id:
+ related_entity = self.entities_by_id[link.target_id]
+ if related_entity not in result:
+ result.append(related_entity)
+
+ if as_target:
+ # Ищем связи, где сущность является целью
+ relations = self.entity_targets.get(entity_id, [])
+
+ for link in relations:
+ if relation_name is None or link.name == relation_name:
+ # Добавляем саму связь
+ if link not in result:
+ result.append(link)
+
+ # Добавляем исходную сущность
+ if link.source_id in self.entities_by_id:
+ related_entity = self.entities_by_id[link.source_id]
+ if related_entity not in result:
+ result.append(related_entity)
+
+ return result
\ No newline at end of file
diff --git a/lib/extractor/ntr_text_fragmentation/core/injection_builder.py b/lib/extractor/ntr_text_fragmentation/core/injection_builder.py
new file mode 100644
index 0000000000000000000000000000000000000000..6374ff1e0693c1dca496134cad31bab4620d2c44
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/core/injection_builder.py
@@ -0,0 +1,429 @@
+"""
+Класс для сборки документа из чанков.
+"""
+
+from collections import defaultdict
+from typing import Optional, Type
+from uuid import UUID
+
+from ..chunking.chunking_strategy import ChunkingStrategy
+from ..models.chunk import Chunk
+from ..models.linker_entity import LinkerEntity
+from .entity_repository import EntityRepository, InMemoryEntityRepository
+
+
+class InjectionBuilder:
+ """
+ Класс для сборки документов из чанков и связей.
+
+ Отвечает за:
+ - Сборку текста из чанков с учетом порядка
+ - Ранжирование документов на основе весов чанков
+ - Добавление соседних чанков для улучшения сборки
+ - Сборку данных из таблиц и других сущностей
+ """
+
+ def __init__(
+ self,
+ repository: EntityRepository | None = None,
+ entities: list[LinkerEntity] | None = None,
+ ):
+ """
+ Инициализация сборщика инъекций.
+
+ Args:
+ repository: Репозиторий сущностей (если None, используется InMemoryEntityRepository)
+ entities: Список всех сущностей (опционально, для обратной совместимости)
+ """
+ # Для обратной совместимости
+ if repository is None and entities is not None:
+ repository = InMemoryEntityRepository(entities)
+
+ self.repository = repository or InMemoryEntityRepository()
+ self.strategy_map: dict[str, Type[ChunkingStrategy]] = {}
+
+ def register_strategy(
+ self,
+ doc_type: str,
+ strategy: Type[ChunkingStrategy],
+ ) -> None:
+ """
+ Регистрирует стратегию для определенного типа документа.
+
+ Args:
+ doc_type: Тип документа
+ strategy: Стратегия чанкинга
+ """
+ self.strategy_map[doc_type] = strategy
+
+ def build(
+ self,
+ filtered_entities: list[LinkerEntity] | list[UUID],
+ chunk_scores: dict[str, float] | None = None,
+ include_tables: bool = True,
+ max_documents: Optional[int] = None,
+ ) -> str:
+ """
+ Собирает текст из всех документов, связанных с предоставленными чанками.
+
+ Args:
+ filtered_entities: Список чанков или их идентификаторов
+ chunk_scores: Словарь весов чанков {chunk_id: score}
+ include_tables: Флаг для включения таблиц в результат
+ max_documents: Максимальное количество документов (None = все)
+
+ Returns:
+ Собранный текст со всеми документами
+ """
+ # Преобразуем входные данные в список идентификаторов
+ entity_ids = [
+ entity.id if isinstance(entity, LinkerEntity) else entity
+ for entity in filtered_entities
+ ]
+
+ print(f"entity_ids: {entity_ids[:3]}...{entity_ids[-3:]}")
+
+ if not entity_ids:
+ return ""
+
+ # Получаем сущности по их идентификаторам
+ entities = self.repository.get_entities_by_ids(entity_ids)
+
+ print(f"entities: {entities[:3]}...{entities[-3:]}")
+
+ # Десериализуем сущности в их специализированные типы
+ deserialized_entities = []
+ for entity in entities:
+ # Используем статический метод десериализации
+ deserialized_entity = LinkerEntity.deserialize(entity)
+ deserialized_entities.append(deserialized_entity)
+
+ print(f"deserialized_entities: {deserialized_entities[:3]}...{deserialized_entities[-3:]}")
+
+ # Фильтруем сущности на чанки и таблицы
+ chunks = [e for e in deserialized_entities if "Chunk" in e.type]
+ tables = [e for e in deserialized_entities if "Table" in e.type]
+
+ # Группируем таблицы по документам
+ table_ids = {table.id for table in tables}
+ doc_tables = self._group_tables_by_document(table_ids)
+
+ if not chunks and not tables:
+ return ""
+
+ # Получаем идентификаторы чанков
+ chunk_ids = [chunk.id for chunk in chunks]
+
+ # Получаем связи для чанков (чанки являются целями связей)
+ links = self.repository.get_related_entities(
+ chunk_ids,
+ relation_name="document_to_chunk",
+ as_target=True,
+ )
+
+ print(f"links: {links[:3]}...{links[-3:]}")
+
+ # Группируем чанки по документам
+ doc_chunks = self._group_chunks_by_document(chunks, links)
+
+ print(f"doc_chunks: {doc_chunks}")
+
+ # Получаем все документы для чанков и таблиц
+ doc_ids = set(doc_chunks.keys()) | set(doc_tables.keys())
+ docs = self.repository.get_entities_by_ids(doc_ids)
+
+ # Десериализуем документы
+ deserialized_docs = []
+ for doc in docs:
+ deserialized_doc = LinkerEntity.deserialize(doc)
+ deserialized_docs.append(deserialized_doc)
+
+ print(f"deserialized_docs: {deserialized_docs[:3]}...{deserialized_docs[-3:]}")
+
+ # Вычисляем веса документов на основе весов чанков
+ doc_scores = self._calculate_document_scores(doc_chunks, chunk_scores)
+
+ # Сортируем документы по весам (по убыванию)
+ sorted_docs = sorted(
+ deserialized_docs,
+ key=lambda d: doc_scores.get(str(d.id), 0.0),
+ reverse=True
+ )
+
+ print(f"sorted_docs: {sorted_docs[:3]}...{sorted_docs[-3:]}")
+
+ # Ограничиваем количество документов, если указано
+ if max_documents:
+ sorted_docs = sorted_docs[:max_documents]
+
+ print(f"sorted_docs: {sorted_docs[:3]}...{sorted_docs[-3:]}")
+
+ # Собираем текст для каждого документа
+ result_parts = []
+ for doc in sorted_docs:
+ doc_text = self._build_document_text(
+ doc,
+ doc_chunks.get(doc.id, []),
+ doc_tables.get(doc.id, []),
+ include_tables
+ )
+ if doc_text:
+ result_parts.append(doc_text)
+
+ # Объединяем результаты
+ return "\n\n".join(result_parts)
+
+ def _build_document_text(
+ self,
+ doc: LinkerEntity,
+ chunks: list[LinkerEntity],
+ tables: list[LinkerEntity],
+ include_tables: bool
+ ) -> str:
+ """
+ Собирает текст документа из чанков и таблиц.
+
+ Args:
+ doc: Сущность документа
+ chunks: Список чанков документа
+ tables: Список таблиц документа
+ include_tables: Флаг для включения таблиц
+
+ Returns:
+ Собранный текст документа
+ """
+ # Получаем стратегию чанкинга
+ strategy_name = doc.metadata.get("chunking_strategy", "fixed_size")
+ strategy = self._get_strategy_instance(strategy_name)
+
+ # Собираем текст из чанков
+ chunks_text = strategy.dechunk(chunks, self.repository) if chunks else ""
+
+ # Собираем текст из таблиц, если нужно
+ tables_text = ""
+ if include_tables and tables:
+ # Сортируем таблицы по индексу, если он есть
+ sorted_tables = sorted(
+ tables,
+ key=lambda t: t.metadata.get("table_index", 0) if t.metadata else 0
+ )
+
+ # Собираем текст таблиц
+ tables_text = "\n\n".join(table.text for table in sorted_tables if hasattr(table, 'text'))
+
+ # Формируем результат
+ result = f"[Источник] - {doc.name}\n"
+ if chunks_text:
+ result += chunks_text
+ if tables_text:
+ if chunks_text:
+ result += "\n\n"
+ result += tables_text
+
+ return result
+
+ def _group_chunks_by_document(
+ self,
+ chunks: list[LinkerEntity],
+ links: list[LinkerEntity]
+ ) -> dict[UUID, list[LinkerEntity]]:
+ """
+ Группирует чанки по документам.
+
+ Args:
+ chunks: Список чанков
+ links: Список связей между документами и чанками
+
+ Returns:
+ Словарь {doc_id: [chunks]}
+ """
+ result = defaultdict(list)
+
+ # Создаем словарь для быстрого доступа к чанкам по ID
+ chunk_dict = {chunk.id: chunk for chunk in chunks}
+
+ # Группируем чанки по документам на основе связей
+ for link in links:
+ if link.target_id in chunk_dict and link.source_id:
+ result[link.source_id].append(chunk_dict[link.target_id])
+
+ return result
+
+ def _group_tables_by_document(
+ self,
+ table_ids: set[UUID]
+ ) -> dict[UUID, list[LinkerEntity]]:
+ """
+ Группирует таблицы по документам.
+
+ Args:
+ table_ids: Множество идентификаторов таблиц
+
+ Returns:
+ Словарь {doc_id: [tables]}
+ """
+ result = defaultdict(list)
+
+ table_ids = [str(table_id) for table_id in table_ids]
+
+ # Получаем связи для таблиц (таблицы являются целями связей)
+ if not table_ids:
+ return result
+
+ links = self.repository.get_related_entities(
+ table_ids,
+ relation_name="document_to_table",
+ as_target=True,
+ )
+
+ # Получаем сами таблицы
+ tables = self.repository.get_entities_by_ids(table_ids)
+
+ # Десериализуем таблицы
+ deserialized_tables = []
+ for table in tables:
+ deserialized_table = LinkerEntity.deserialize(table)
+ deserialized_tables.append(deserialized_table)
+
+ # Создаем словарь для быстрого доступа к таблицам по ID
+ table_dict = {str(table.id): table for table in deserialized_tables}
+
+ # Группируем таблицы по документам на основе связей
+ for link in links:
+ if link.target_id in table_dict and link.source_id:
+ result[link.source_id].append(table_dict[link.target_id])
+
+ return result
+
+ def _calculate_document_scores(
+ self,
+ doc_chunks: dict[UUID, list[LinkerEntity]],
+ chunk_scores: Optional[dict[str, float]]
+ ) -> dict[str, float]:
+ """
+ Вычисляет веса документов на основе весов чанков.
+
+ Args:
+ doc_chunks: Словарь {doc_id: [chunks]}
+ chunk_scores: Словарь весов чанков {chunk_id: score}
+
+ Returns:
+ Словарь весов документов {doc_id: score}
+ """
+ if not chunk_scores:
+ return {str(doc_id): 1.0 for doc_id in doc_chunks.keys()}
+
+ result = {}
+ for doc_id, chunks in doc_chunks.items():
+ # Берем максимальный вес среди чанков документа
+ chunk_weights = [chunk_scores.get(str(c.id), 0.0) for c in chunks]
+ result[str(doc_id)] = max(chunk_weights) if chunk_weights else 0.0
+
+ return result
+
+ def add_neighboring_chunks(
+ self, entities: list[LinkerEntity] | list[UUID], max_distance: int = 1
+ ) -> list[LinkerEntity]:
+ """
+ Добавляет соседние чанки к отфильтрованному списку чанков.
+
+ Args:
+ entities: Список сущностей или их идентификаторов
+ max_distance: Максимальное расстояние для поиска соседей
+
+ Returns:
+ Расширенный список сущностей
+ """
+ # Преобразуем входные данные в список идентификаторов
+ entity_ids = [
+ entity.id if isinstance(entity, LinkerEntity) else entity
+ for entity in entities
+ ]
+
+ if not entity_ids:
+ return []
+
+ # Получаем исходные сущности
+ original_entities = self.repository.get_entities_by_ids(entity_ids)
+
+ # Фильтруем только чанки
+ chunk_entities = [e for e in original_entities if isinstance(e, Chunk)]
+
+ if not chunk_entities:
+ return original_entities
+
+ # Получаем идентификаторы чанков
+ chunk_ids = [chunk.id for chunk in chunk_entities]
+
+ # Получаем соседние чанки
+ neighboring_chunks = self.repository.get_neighboring_chunks(
+ chunk_ids, max_distance
+ )
+
+ # Объединяем исходные сущности с соседними чанками
+ result = list(original_entities)
+ for chunk in neighboring_chunks:
+ if chunk not in result:
+ result.append(chunk)
+
+ # Получаем документы и связи для всех чанков
+ all_chunk_ids = [chunk.id for chunk in result if isinstance(chunk, Chunk)]
+
+ docs = self.repository.get_document_for_chunks(all_chunk_ids)
+ links = self.repository.get_related_entities(
+ all_chunk_ids, relation_name="document_to_chunk", as_target=True
+ )
+
+ # Добавляем документы и связи в результат
+ for doc in docs:
+ if doc not in result:
+ result.append(doc)
+
+ for link in links:
+ if link not in result:
+ result.append(link)
+
+ return result
+
+ def _get_strategy_instance(self, strategy_name: str) -> ChunkingStrategy:
+ """
+ Создает экземпляр стратегии чанкинга по имени.
+
+ Args:
+ strategy_name: Имя стратегии
+
+ Returns:
+ Экземпляр соответствующей стратегии
+ """
+ # Используем словарь для маппинга имен стратегий на их классы
+ strategies = {
+ "fixed_size": "..chunking.specific_strategies.fixed_size_chunking.FixedSizeChunkingStrategy",
+ }
+
+ # Если стратегия зарегистрирована в self.strategy_map, используем её
+ if strategy_name in self.strategy_map:
+ return self.strategy_map[strategy_name]()
+
+ # Если стратегия известна, импортируем и инициализируем её
+ if strategy_name in strategies:
+ import importlib
+
+ module_path, class_name = strategies[strategy_name].rsplit(".", 1)
+ try:
+ # Конвертируем относительный путь в абсолютный
+ abs_module_path = f"ntr_text_fragmentation{module_path[2:]}"
+ module = importlib.import_module(abs_module_path)
+ strategy_class = getattr(module, class_name)
+ return strategy_class()
+ except (ImportError, AttributeError) as e:
+ # Если импорт не удался, используем стратегию по умолчанию
+ from ..chunking.specific_strategies.fixed_size_chunking import \
+ FixedSizeChunkingStrategy
+
+ return FixedSizeChunkingStrategy()
+
+ # По умолчанию используем стратегию с фиксированным размером
+ from ..chunking.specific_strategies.fixed_size_chunking import \
+ FixedSizeChunkingStrategy
+
+ return FixedSizeChunkingStrategy()
diff --git a/lib/extractor/ntr_text_fragmentation/integrations/__init__.py b/lib/extractor/ntr_text_fragmentation/integrations/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a615b54c85be9b31d1b45546c4e3314314e49bd7
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/integrations/__init__.py
@@ -0,0 +1,9 @@
+"""
+Модуль интеграций с внешними хранилищами данных и ORM системами.
+"""
+
+from .sqlalchemy_repository import SQLAlchemyEntityRepository
+
+__all__ = [
+ "SQLAlchemyEntityRepository",
+]
diff --git a/lib/extractor/ntr_text_fragmentation/integrations/sqlalchemy_repository.py b/lib/extractor/ntr_text_fragmentation/integrations/sqlalchemy_repository.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f20cce31176305a0b0b2c5e3d7523f35c20c197
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/integrations/sqlalchemy_repository.py
@@ -0,0 +1,339 @@
+"""
+Реализация EntityRepository для работы с SQLAlchemy.
+"""
+
+from abc import abstractmethod
+from typing import Any, Iterable, List, Optional, Type
+from uuid import UUID
+
+from sqlalchemy import and_, select
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import Session
+
+from ..core.entity_repository import EntityRepository
+from ..models import Chunk, LinkerEntity
+
+Base = declarative_base()
+
+
+class SQLAlchemyEntityRepository(EntityRepository):
+ """
+ Реализация EntityRepository для работы с базой данных через SQLAlchemy.
+
+ Эта реализация предполагает, что таблицы для хранения сущностей уже созданы
+ в базе данных и соответствуют определенной структуре.
+
+ Вы можете наследоваться от этого класса и определить свою структуру моделей,
+ переопределив абстрактные методы.
+ """
+
+ def __init__(self, db: Session):
+ """
+ Инициализирует репозиторий с указанной сессией SQLAlchemy.
+
+ Args:
+ db: Сессия SQLAlchemy для работы с базой данных
+ """
+ self.db = db
+
+ @abstractmethod
+ def _entity_model_class(self) -> Type['Base']:
+ """
+ Возвращает класс модели SQLAlchemy для сущностей.
+
+ Returns:
+ Класс модели SQLAlchemy для сущностей
+ """
+ pass
+
+ @abstractmethod
+ def _map_db_entity_to_linker_entity(self, db_entity: Any) -> LinkerEntity:
+ """
+ Преобразует сущность из базы данных в LinkerEntity.
+
+ Args:
+ db_entity: Сущность из базы данных
+
+ Returns:
+ Сущность LinkerEntity
+ """
+ pass
+
+ def get_entities_by_ids(self, entity_ids: Iterable[UUID]) -> List[LinkerEntity]:
+ """
+ Получить сущности по списку идентификаторов.
+
+ Args:
+ entity_ids: Список идентификаторов сущностей
+
+ Returns:
+ Список сущностей, соответствующих указанным идентификаторам
+ """
+ if not entity_ids:
+ return []
+
+ with self.db() as session:
+ entity_model = self._entity_model_class()
+ db_entities = session.execute(
+ select(entity_model).where(entity_model.uuid.in_(list(entity_ids)))
+ ).scalars().all()
+ print(f"db_entities: {db_entities[:3]}...{db_entities[-3:]}")
+
+ mapped_entities = [self._map_db_entity_to_linker_entity(entity) for entity in db_entities]
+ print(f"mapped_entities: {mapped_entities[:3]}...{mapped_entities[-3:]}")
+ return mapped_entities
+
+ def get_document_for_chunks(self, chunk_ids: Iterable[UUID]) -> List[LinkerEntity]:
+ """
+ Получить документы, которым принадлежат указанные чанки.
+
+ Args:
+ chunk_ids: Список идентификаторов чанков
+
+ Returns:
+ Список документов, которым принадлежат указанные чанки
+ """
+ if not chunk_ids:
+ return []
+
+ with self.db() as session:
+ entity_model = self._entity_model_class()
+
+ string_ids = [str(id) for id in chunk_ids]
+
+ # Получаем все сущности-связи между документами и чанками
+ links = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.target_id.in_(string_ids),
+ entity_model.name == "document_to_chunk",
+ entity_model.target_id.isnot(None) # Проверяем, что это связь
+ )
+ )
+ ).scalars().all()
+
+ if not links:
+ return []
+
+ # Извлекаем ID документов
+ doc_ids = [link.source_id for link in links]
+
+ # Получаем документы по их ID
+ documents = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.uuid.in_(doc_ids),
+ entity_model.entity_type == "DocumentAsEntity"
+ )
+ )
+ ).scalars().all()
+
+ return [self._map_db_entity_to_linker_entity(doc) for doc in documents]
+
+ def get_neighboring_chunks(self,
+ chunk_ids: Iterable[UUID],
+ max_distance: int = 1) -> List[LinkerEntity]:
+ """
+ Получить соседние чанки для указанных чанков.
+
+ Args:
+ chunk_ids: Список идентификаторов чанков
+ max_distance: Максимальное расстояние до соседа
+
+ Returns:
+ Список соседних чанков
+ """
+ if not chunk_ids:
+ return []
+
+ string_ids = [str(id) for id in chunk_ids]
+
+ with self.db() as session:
+ entity_model = self._entity_model_class()
+ result = []
+
+ # Сначала получаем указанные чанки, чтобы узнать их индексы и документы
+ chunks = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.uuid.in_(string_ids),
+ entity_model.entity_type.like("%Chunk") # Используем LIKE для поиска всех типов чанков
+ )
+ )
+ ).scalars().all()
+
+ print(f"chunks: {chunks[:3]}...{chunks[-3:]}")
+
+ if not chunks:
+ return []
+
+ # Находим документы для чанков через связи
+ doc_ids = set()
+ chunk_indices = {}
+
+ for chunk in chunks:
+ mapped_chunk = self._map_db_entity_to_linker_entity(chunk)
+ if not isinstance(mapped_chunk, Chunk):
+ continue
+
+ chunk_indices[chunk.uuid] = mapped_chunk.chunk_index
+
+ # Находим связь от документа к чанку
+ links = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.target_id == chunk.uuid,
+ entity_model.name == "document_to_chunk"
+ )
+ )
+ ).scalars().all()
+
+ print(f"links: {links[:3]}...{links[-3:]}")
+
+ for link in links:
+ doc_ids.add(link.source_id)
+
+ if not doc_ids or not any(idx is not None for idx in chunk_indices.values()):
+ return []
+
+ # Для каждого документа находим все его чанки
+ for doc_id in doc_ids:
+ # Находим все связи от документа к чанкам
+ links = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.source_id == doc_id,
+ entity_model.name == "document_to_chunk"
+ )
+ )
+ ).scalars().all()
+
+ doc_chunk_ids = [link.target_id for link in links]
+
+ print(f"doc_chunk_ids: {doc_chunk_ids[:3]}...{doc_chunk_ids[-3:]}")
+
+ # Получаем все чанки документа
+ doc_chunks = session.execute(
+ select(entity_model).where(
+ and_(
+ entity_model.uuid.in_(doc_chunk_ids),
+ entity_model.entity_type.like("%Chunk") # Используем LIKE для поиска всех типов чанков
+ )
+ )
+ ).scalars().all()
+
+ print(f"doc_chunks: {doc_chunks[:3]}...{doc_chunks[-3:]}")
+
+ # Для каждого чанка в документе проверяем, является ли он соседом
+ for doc_chunk in doc_chunks:
+ if doc_chunk.uuid in chunk_ids:
+ continue
+
+ mapped_chunk = self._map_db_entity_to_linker_entity(doc_chunk)
+ if not isinstance(mapped_chunk, Chunk):
+ continue
+
+ chunk_index = mapped_chunk.chunk_index
+ if chunk_index is None:
+ continue
+
+ # Проверяем, является ли чанк соседом какого-либо из исходных чанков
+ is_neighbor = False
+ for orig_chunk_id, orig_index in chunk_indices.items():
+ if orig_index is not None and abs(chunk_index - orig_index) <= max_distance:
+ is_neighbor = True
+ break
+
+ if is_neighbor:
+ result.append(mapped_chunk)
+
+ return result
+
+ def get_related_entities(self,
+ entity_ids: Iterable[UUID],
+ relation_name: Optional[str] = None,
+ as_source: bool = False,
+ as_target: bool = False) -> List[LinkerEntity]:
+ """
+ Получить сущности, связанные с указанными сущностями.
+
+ Args:
+ entity_ids: Список идентификаторов сущностей
+ relation_name: Опциональное имя отношения для фильтрации
+ as_source: Если True, ищем связи, где указанные entity_ids являются источниками
+ as_target: Если True, ищем связи, где указанные entity_ids являются целями
+
+ Returns:
+ Список связанных сущностей и связей
+ """
+ if not entity_ids:
+ return []
+
+ entity_model = self._entity_model_class()
+ result = []
+
+ # Если не указано ни as_source, ни as_target, по умолчанию ищем связи,
+ # где указанные entity_ids являются источниками
+ if not as_source and not as_target:
+ as_source = True
+
+ string_ids = [str(id) for id in entity_ids]
+
+ with self.db() as session:
+ # Поиск связей, где указанные entity_ids являются источниками
+ if as_source:
+ conditions = [
+ entity_model.source_id.in_(string_ids)
+ ]
+
+ if relation_name:
+ conditions.append(entity_model.name == relation_name)
+
+ links = session.execute(
+ select(entity_model).where(and_(*conditions))
+ ).scalars().all()
+
+ for link in links:
+ # Добавляем связь
+ link_entity = self._map_db_entity_to_linker_entity(link)
+ result.append(link_entity)
+
+ # Добавляем целевую сущность
+ target_entities = session.execute(
+ select(entity_model).where(entity_model.uuid == link.target_id)
+ ).scalars().all()
+
+ for target in target_entities:
+ target_entity = self._map_db_entity_to_linker_entity(target)
+ if target_entity not in result:
+ result.append(target_entity)
+
+ # Поиск связей, где указанные entity_ids являются целями
+ if as_target:
+ conditions = [
+ entity_model.target_id.in_(string_ids)
+ ]
+
+ if relation_name:
+ conditions.append(entity_model.name == relation_name)
+
+ links = session.execute(
+ select(entity_model).where(and_(*conditions))
+ ).scalars().all()
+
+ for link in links:
+ # Добавляем связь
+ link_entity = self._map_db_entity_to_linker_entity(link)
+ result.append(link_entity)
+
+ # Добавляем исходную сущность
+ source_entities = session.execute(
+ select(entity_model).where(entity_model.uuid == link.source_id)
+ ).scalars().all()
+
+ for source in source_entities:
+ source_entity = self._map_db_entity_to_linker_entity(source)
+ if source_entity not in result:
+ result.append(source_entity)
+
+ return result
diff --git a/lib/extractor/ntr_text_fragmentation/models/__init__.py b/lib/extractor/ntr_text_fragmentation/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..738b25982483c288bf71d1ab863557a9e37ee654
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/models/__init__.py
@@ -0,0 +1,13 @@
+"""
+Модуль моделей данных.
+"""
+
+from .chunk import Chunk
+from .document import DocumentAsEntity
+from .linker_entity import LinkerEntity
+
+__all__ = [
+ "LinkerEntity",
+ "DocumentAsEntity",
+ "Chunk",
+]
\ No newline at end of file
diff --git a/lib/extractor/ntr_text_fragmentation/models/chunk.py b/lib/extractor/ntr_text_fragmentation/models/chunk.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a054343ddd1ce0935438d3805d17b92ac9d7229
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/models/chunk.py
@@ -0,0 +1,48 @@
+"""
+Класс для представления чанка документа.
+"""
+
+from dataclasses import dataclass
+
+from .linker_entity import LinkerEntity, register_entity
+
+
+@register_entity
+@dataclass
+class Chunk(LinkerEntity):
+ """
+ Класс для представления чанка документа в системе извлечения и сборки.
+
+ Attributes:
+ chunk_index: Порядковый номер чанка в документе (0-based).
+ Используется для восстановления порядка при сборке.
+ """
+
+ chunk_index: int | None = None
+
+ @classmethod
+ def deserialize(cls, data: LinkerEntity) -> 'Chunk':
+ """
+ Десериализует Chunk из объекта LinkerEntity.
+
+ Базовый класс Chunk не должен использоваться напрямую,
+ все конкретные реализации должны переопределить этот метод.
+
+ Args:
+ data: Объект LinkerEntity для преобразования в Chunk
+
+ Raises:
+ NotImplementedError: Метод должен быть переопределен в дочерних классах
+ """
+ if cls == Chunk:
+ # Если это прямой вызов на базовом классе Chunk, выбрасываем исключение
+ raise NotImplementedError(
+ "Базовый класс Chunk не поддерживает десериализацию. "
+ "Используйте конкретную реализацию Chunk (например, FixedSizeChunk)."
+ )
+
+ # Если вызывается из дочернего класса, который не переопределил метод,
+ # выбрасываем более конкретную ошибку
+ raise NotImplementedError(
+ f"Класс {cls.__name__} должен реализовать метод deserialize."
+ )
diff --git a/lib/extractor/ntr_text_fragmentation/models/document.py b/lib/extractor/ntr_text_fragmentation/models/document.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d661a195b938d8ef99a959dec0b86b5afb259f1
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/models/document.py
@@ -0,0 +1,49 @@
+"""
+Класс для представления документа как сущности.
+"""
+
+from dataclasses import dataclass
+
+from .linker_entity import LinkerEntity, register_entity
+
+
+@register_entity
+@dataclass
+class DocumentAsEntity(LinkerEntity):
+ """
+ Класс для представления документа как сущности в системе извлечения и сборки.
+ """
+
+ doc_type: str = "unknown"
+
+ @classmethod
+ def deserialize(cls, data: LinkerEntity) -> 'DocumentAsEntity':
+ """
+ Десериализует DocumentAsEntity из объекта LinkerEntity.
+
+ Args:
+ data: Объект LinkerEntity для преобразования в DocumentAsEntity
+
+ Returns:
+ Десериализованный объект DocumentAsEntity
+ """
+ metadata = data.metadata or {}
+
+ # Получаем тип документа из метаданных или используем значение по умолчанию
+ doc_type = metadata.get('_doc_type', 'unknown')
+
+ # Создаем чистые метаданные без служебных полей
+ clean_metadata = {k: v for k, v in metadata.items() if not k.startswith('_')}
+
+ return cls(
+ id=data.id,
+ name=data.name,
+ text=data.text,
+ in_search_text=data.in_search_text,
+ metadata=clean_metadata,
+ source_id=data.source_id,
+ target_id=data.target_id,
+ number_in_relation=data.number_in_relation,
+ type="DocumentAsEntity",
+ doc_type=doc_type
+ )
diff --git a/lib/extractor/ntr_text_fragmentation/models/linker_entity.py b/lib/extractor/ntr_text_fragmentation/models/linker_entity.py
new file mode 100644
index 0000000000000000000000000000000000000000..05619db99e101ffbbd9b4ccc48dd46b6918ffa63
--- /dev/null
+++ b/lib/extractor/ntr_text_fragmentation/models/linker_entity.py
@@ -0,0 +1,217 @@
+"""
+Базовый абстрактный класс для всех сущностей с поддержкой триплетного подхода.
+"""
+
+import uuid
+from abc import abstractmethod
+from dataclasses import dataclass, field, fields
+from uuid import UUID
+
+
+@dataclass
+class LinkerEntity:
+ """
+ Общий класс для всех сущностей в системе извлечения и сборки.
+ Поддерживает триплетный подход, где каждая сущность может опционально связывать две другие сущности.
+
+ Attributes:
+ id (UUID): Уникальный идентификатор сущности.
+ name (str): Название сущности.
+ text (str): Текстое представление сущности.
+ in_search_text (str | None): Текст для поиска. Если задан, используется в __str__, иначе используется обычное представление.
+ metadata (dict): Метаданные сущности.
+ source_id (UUID | None): Опциональный идентификатор исходной сущности.
+ Если указан, эта сущность является связью.
+ target_id (UUID | None): Опциональный идентификатор целевой сущности.
+ Если указан, эта сущность является связью.
+ number_in_relation (int | None): Используется в случае связей один-ко-многим,
+ указывает номер целевой сущности в списке.
+ type (str): Тип сущности.
+ """
+
+ id: UUID
+ name: str
+ text: str
+ metadata: dict # JSON с метаданными
+ in_search_text: str | None = None
+ source_id: UUID | None = None
+ target_id: UUID | None = None
+ number_in_relation: int | None = None
+ type: str = field(default_factory=lambda: "Entity")
+
+ def __post_init__(self):
+ if self.id is None:
+ self.id = uuid.uuid4()
+
+ # Проверяем корректность полей связи
+ if (self.source_id is not None and self.target_id is None) or \
+ (self.source_id is None and self.target_id is not None):
+ raise ValueError("source_id и target_id должны быть либо оба указаны, либо оба None")
+
+ def is_link(self) -> bool:
+ """
+ Проверяет, является ли сущность связью (имеет и source_id, и target_id).
+
+ Returns:
+ bool: True, если сущность является связью, иначе False
+ """
+ return self.source_id is not None and self.target_id is not None
+
+ def __str__(self) -> str:
+ """
+ Возвращает строковое представление сущности.
+ Если задан in_search_text, возвращает его, иначе возвращает стандартное представление.
+ """
+ if self.in_search_text is not None:
+ return self.in_search_text
+ return f"{self.name}: {self.text}"
+
+ def __eq__(self, other: 'LinkerEntity') -> bool:
+ """
+ Сравнивает текущую сущность с другой.
+
+ Args:
+ other: Другая сущность для сравнения
+
+ Returns:
+ bool: True если сущности совпадают, иначе False
+ """
+ if not isinstance(other, self.__class__):
+ return False
+
+ basic_equality = (
+ self.id == other.id
+ and self.name == other.name
+ and self.text == other.text
+ and self.type == other.type
+ )
+
+ # Если мы имеем дело со связями, также проверяем поля связи
+ if self.is_link() or other.is_link():
+ return (
+ basic_equality
+ and self.source_id == other.source_id
+ and self.target_id == other.target_id
+ )
+
+ return basic_equality
+
+ def serialize(self) -> 'LinkerEntity':
+ """
+ Сериализует сущность в простейшую форму сущности, передавая все дополнительные поля в метаданные.
+ """
+ # Получаем список полей базового класса
+ known_fields = {field.name for field in fields(LinkerEntity)}
+
+ # Получаем все атрибуты текущего объекта
+ dict_entity = {}
+ for attr_name in dir(self):
+ # Пропускаем служебные атрибуты, методы и уже известные поля
+ if (
+ attr_name.startswith('_')
+ or attr_name in known_fields
+ or callable(getattr(self, attr_name))
+ ):
+ continue
+
+ # Добавляем дополнительные поля в словарь
+ dict_entity[attr_name] = getattr(self, attr_name)
+
+ # Преобразуем имена дополнительных полей, добавляя префикс "_"
+ dict_entity = {f'_{name}': value for name, value in dict_entity.items()}
+
+ # Объединяем с существующими метаданными
+ dict_entity = {**dict_entity, **self.metadata}
+
+ result_type = self.type
+ if result_type == "Entity":
+ result_type = self.__class__.__name__
+
+ # Создаем базовый объект LinkerEntity с новыми метаданными
+ return LinkerEntity(
+ id=self.id,
+ name=self.name,
+ text=self.text,
+ in_search_text=self.in_search_text,
+ metadata=dict_entity,
+ source_id=self.source_id,
+ target_id=self.target_id,
+ number_in_relation=self.number_in_relation,
+ type=result_type,
+ )
+
+ @classmethod
+ @abstractmethod
+ def deserialize(cls, data: 'LinkerEntity') -> 'Self':
+ """
+ Десериализует сущность из простейшей формы сущности, учитывая все дополнительные поля в метаданных.
+ """
+ raise NotImplementedError(
+ f"Метод deserialize для класса {cls.__class__.__name__} не реализован"
+ )
+
+ # Реестр для хранения всех наследников LinkerEntity
+ _entity_classes = {}
+
+ @classmethod
+ def register_entity_class(cls, entity_class):
+ """
+ Регистрирует класс-наследник в реестре.
+
+ Args:
+ entity_class: Класс для регистрации
+ """
+ entity_type = entity_class.__name__
+ cls._entity_classes[entity_type] = entity_class
+ # Также регистрируем по типу, если он отличается от имени класса
+ if hasattr(entity_class, 'type') and isinstance(entity_class.type, str):
+ cls._entity_classes[entity_class.type] = entity_class
+
+ @classmethod
+ def deserialize(cls, data: 'LinkerEntity') -> 'LinkerEntity':
+ """
+ Десериализует сущность в нужный тип на основе поля type.
+
+ Args:
+ data: Сериализованная сущность типа LinkerEntity
+
+ Returns:
+ Десериализованная сущность правильного типа
+ """
+ # Получаем тип сущности
+ entity_type = data.type
+
+ # Проверяем реестр классов
+ if entity_type in cls._entity_classes:
+ try:
+ return cls._entity_classes[entity_type].deserialize(data)
+ except (AttributeError, NotImplementedError) as e:
+ # Если метод не реализован, возвращаем исходную сущность
+ return data
+
+ # Если тип не найден в реестре, просто возвращаем исходную сущность
+ # Больше не используем опасное сканирование sys.modules
+ return data
+
+
+# Декоратор для регистрации производных классов
+def register_entity(cls):
+ """
+ Декоратор для регистрации классов-наследников LinkerEntity.
+
+ Пример использования:
+
+ @register_entity
+ class MyEntity(LinkerEntity):
+ type = "my_entity"
+
+ Args:
+ cls: Класс, который нужно зарегистрировать
+
+ Returns:
+ Исходный класс (без изменений)
+ """
+ # Регистрируем класс в реестр, используя его имя или указанный тип
+ entity_type = getattr(cls, 'type', cls.__name__)
+ LinkerEntity._entity_classes[entity_type] = cls
+ return cls
diff --git a/lib/extractor/pyproject.toml b/lib/extractor/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..a9a134fb7f2eb3b778d904c355298f6d4499821c
--- /dev/null
+++ b/lib/extractor/pyproject.toml
@@ -0,0 +1,26 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = ["setuptools>=61"]
+
+[project]
+name = "ntr_text_fragmentation"
+version = "0.1.0"
+dependencies = [
+ "uuid==1.30",
+ "ntr_fileparser @ git+ssh://git@gitlab.ntrlab.ru/textai/parsers/parser.git@master"
+]
+
+[project.optional-dependencies]
+test = [
+ "pytest>=7.0.0",
+ "pytest-cov>=4.0.0"
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+
+[tool.pytest]
+testpaths = ["tests"]
+python_files = "test_*.py"
+python_classes = "Test*"
+python_functions = "test_*"
\ No newline at end of file
diff --git a/lib/extractor/scripts/README_test_chunking.md b/lib/extractor/scripts/README_test_chunking.md
new file mode 100644
index 0000000000000000000000000000000000000000..cf31de3049819dc490aa8f6a2da5bb2136f17e6e
--- /dev/null
+++ b/lib/extractor/scripts/README_test_chunking.md
@@ -0,0 +1,107 @@
+# Тестирование чанкинга и сборки документов
+
+Скрипт `test_chunking.py` позволяет тестировать различные стратегии чанкинга документов и их последующую сборку.
+
+## Возможности
+
+1. **Разбивка документов** - применение различных стратегий чанкинга к документам
+2. **Сохранение результатов** - сохранение чанков и метаданных в CSV
+3. **Сборка документов** - загрузка чанков из CSV и сборка документа с помощью InjectionBuilder
+4. **Фильтрация чанков** - возможность фильтровать чанки по индексу или ключевым словам
+
+## Режимы работы
+
+Скрипт поддерживает три режима работы:
+
+1. **chunk** - только разбивка документа на чанки и сохранение в CSV
+2. **build** - загрузка чанков из CSV и сборка документа
+3. **full** - разбивка документа, сохранение в CSV и последующая сборка
+
+## Примеры использования
+
+### Разбивка документа на чанки (стратегия fixed_size)
+
+```bash
+python scripts/test_chunking.py --mode chunk --input test_input/test.docx --strategy fixed_size --words 50 --overlap 25
+```
+
+### Разбивка документа на чанки (стратегия sentence)
+
+```bash
+python scripts/test_chunking.py --mode chunk --input test_input/test.docx --strategy sentence
+```
+
+### Загрузка чанков из CSV и сборка документа (все чанки)
+
+```bash
+python scripts/test_chunking.py --mode build --csv test_output/test_fixed_size_w50_o25.csv
+```
+
+### Загрузка чанков из CSV и сборка документа (с фильтрацией по индексу)
+
+```bash
+python scripts/test_chunking.py --mode build --csv test_output/test_fixed_size_w50_o25.csv --filter index --filter-value "0,2,4"
+```
+
+### Загрузка чанков из CSV и сборка документа (с фильтрацией по ключевому слову)
+
+```bash
+python scripts/test_chunking.py --mode build --csv test_output/test_fixed_size_w50_o25.csv --filter keyword --filter-value "важно"
+```
+
+### Полный цикл: разбивка, сохранение и сборка
+
+```bash
+python scripts/test_chunking.py --mode full --input test_input/test.docx --strategy fixed_size --words 50 --overlap 25
+```
+
+## Параметры командной строки
+
+### Основные параметры
+
+| Параметр | Описание | Значения по умолчанию |
+|----------|----------|------------------------|
+| `--mode` | Режим работы | `chunk` |
+| `--input` | Путь к входному файлу | `test_input/test.docx` |
+| `--csv` | Путь к CSV файлу с сущностями | None |
+| `--output-dir` | Директория для выходных файлов | `test_output` |
+
+### Параметры стратегии чанкинга
+
+| Параметр | Описание | Значения по умолчанию |
+|----------|----------|------------------------|
+| `--strategy` | Стратегия чанкинга | `fixed_size` |
+| `--words` | Количество слов в чанке (для fixed_size) | 50 |
+| `--overlap` | Перекрытие в словах (для fixed_size) | 25 |
+| `--debug` | Режим отладки (для numbered_items) | False |
+
+### Параметры фильтрации
+
+| Параметр | Описание | Значения по умолчанию |
+|----------|----------|------------------------|
+| `--filter` | Тип фильтрации чанков | `none` |
+| `--filter-value` | Значение для фильтрации | None |
+
+## Подготовка тестовых данных
+
+Для тестирования скрипта вам понадобится документ в формате docx, txt, pdf или другом поддерживаемом формате. Поместите тестовый документ в папку `test_input`.
+
+## Результаты работы
+
+После выполнения скрипта в папке `test_output` будут созданы следующие файлы:
+
+1. **test_{strategy}_....csv** - CSV файл с сущностями (документ, чанки, связи)
+2. **rebuilt_document_{filter}_{filter_value}.txt** - собранный текст документа (при использовании режимов build или full)
+
+## Примечания
+
+- Для различных стратегий чанкинга доступны разные параметры
+- При сборке документа можно использовать фильтрацию чанков по индексу или ключевому слову
+- Собранный документ будет отличаться от исходного, если использовалась фильтрация чанков
+
+## Требования
+
+- Python 3.8+
+- pandas
+- ntr_fileparser
+- ntr_text_fragmentation
\ No newline at end of file
diff --git a/lib/extractor/scripts/analyze_missing_puncts.py b/lib/extractor/scripts/analyze_missing_puncts.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d381932227953ff9ba1fb79606d2f51b36f3472
--- /dev/null
+++ b/lib/extractor/scripts/analyze_missing_puncts.py
@@ -0,0 +1,547 @@
+#!/usr/bin/env python
+"""
+Скрипт для анализа ненайденных пунктов по лучшему подходу чанкинга (200 слов, 75 перекрытие, baai/bge-m3, top-100).
+Формирует отчет в формате Markdown с топ-5 наиболее похожими чанками для каждого ненайденного пункта.
+"""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from fuzzywuzzy import fuzz
+from sklearn.metrics.pairwise import cosine_similarity
+from tqdm import tqdm
+
+# Константы
+DATA_FOLDER = "data/docs" # Путь к папке с документами
+MODEL_NAME = "BAAI/bge-m3" # Название лучшей модели
+DATASET_PATH = "data/dataset.xlsx" # Путь к Excel-датасету с вопросами
+OUTPUT_DIR = "data" # Директория для сохранения результатов
+MARKDOWN_FILE = "missing_puncts_analysis.md" # Имя выходного MD-файла
+SIMILARITY_THRESHOLD = 0.7 # Порог для нечеткого сравнения
+WORDS_PER_CHUNK = 200 # Размер чанка в словах
+OVERLAP_WORDS = 75 # Перекрытие в словах
+TOP_N = 100 # Количество чанков в топе
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+
+def parse_args():
+ """
+ Парсит аргументы командной строки.
+
+ Returns:
+ Аргументы командной строки
+ """
+ parser = argparse.ArgumentParser(description="Анализ ненайденных пунктов для лучшего подхода чанкинга")
+
+ parser.add_argument("--data-folder", type=str, default=DATA_FOLDER,
+ help=f"Путь к папке с документами (по умолчанию: {DATA_FOLDER})")
+ parser.add_argument("--model-name", type=str, default=MODEL_NAME,
+ help=f"Название модели (по умолчанию: {MODEL_NAME})")
+ parser.add_argument("--dataset-path", type=str, default=DATASET_PATH,
+ help=f"Путь к Excel-датасету с вопросами (по умолчанию: {DATASET_PATH})")
+ parser.add_argument("--output-dir", type=str, default=OUTPUT_DIR,
+ help=f"Директория для сохранения результатов (по умолчанию: {OUTPUT_DIR})")
+ parser.add_argument("--markdown-file", type=str, default=MARKDOWN_FILE,
+ help=f"Имя выходного MD-файла (по умолчанию: {MARKDOWN_FILE})")
+ parser.add_argument("--similarity-threshold", type=float, default=SIMILARITY_THRESHOLD,
+ help=f"Порог для нечеткого сравнения (по умолчанию: {SIMILARITY_THRESHOLD})")
+ parser.add_argument("--words-per-chunk", type=int, default=WORDS_PER_CHUNK,
+ help=f"Размер чанка в словах (по умолчанию: {WORDS_PER_CHUNK})")
+ parser.add_argument("--overlap-words", type=int, default=OVERLAP_WORDS,
+ help=f"Перекрытие в словах (по умолчанию: {OVERLAP_WORDS})")
+ parser.add_argument("--top-n", type=int, default=TOP_N,
+ help=f"Количество чанков в топе (по умолчанию: {TOP_N})")
+
+ return parser.parse_args()
+
+
+def load_questions_dataset(file_path: str) -> pd.DataFrame:
+ """
+ Загружает датасет с вопросами из Excel-файла.
+
+ Args:
+ file_path: Путь к Excel-файлу
+
+ Returns:
+ DataFrame с вопросами и пунктами
+ """
+ print(f"Загрузка датасета из {file_path}...")
+
+ df = pd.read_excel(file_path)
+ print(f"Загружен датасет со столбцами: {df.columns.tolist()}")
+
+ # Преобразуем NaN в пустые строки для текстовых полей
+ text_columns = ['question', 'text', 'item_type']
+ for col in text_columns:
+ if col in df.columns:
+ df[col] = df[col].fillna('')
+
+ return df
+
+
+def load_chunks_and_embeddings(output_dir: str, words_per_chunk: int, overlap_words: int, model_name: str) -> tuple:
+ """
+ Загружает чанки и эмбеддинги из файлов.
+
+ Args:
+ output_dir: Директория с файлами
+ words_per_chunk: Размер чанка в словах
+ overlap_words: Перекрытие в словах
+ model_name: Название модели
+
+ Returns:
+ Кортеж (чанки, эмбеддинги чанков, эмбеддинги вопросов, данные вопросов)
+ """
+ # Формируем уникальное имя для файлов на основе параметров
+ model_name_safe = model_name.replace('/', '_')
+ strategy_config_str = f"fixed_size_w{words_per_chunk}_o{overlap_words}"
+ chunks_filename = f"chunks_{strategy_config_str}_{model_name_safe}"
+ questions_filename = f"questions_{model_name_safe}"
+
+ # Пути к файлам
+ chunks_embeddings_path = os.path.join(output_dir, f"{chunks_filename}_embeddings.npy")
+ chunks_data_path = os.path.join(output_dir, f"{chunks_filename}_data.csv")
+ questions_embeddings_path = os.path.join(output_dir, f"{questions_filename}_embeddings.npy")
+ questions_data_path = os.path.join(output_dir, f"{questions_filename}_data.csv")
+
+ # Проверяем наличие всех файлов
+ for path in [chunks_embeddings_path, chunks_data_path, questions_embeddings_path, questions_data_path]:
+ if not os.path.exists(path):
+ raise FileNotFoundError(f"Файл {path} не найден")
+
+ # Загружаем данные
+ print(f"Загрузка данных из {output_dir}...")
+ chunks_embeddings = np.load(chunks_embeddings_path)
+ chunks_df = pd.read_csv(chunks_data_path)
+ questions_embeddings = np.load(questions_embeddings_path)
+ questions_df = pd.read_csv(questions_data_path)
+
+ print(f"Загружено {len(chunks_df)} чанков и {len(questions_df)} вопросов")
+
+ return chunks_df, chunks_embeddings, questions_embeddings, questions_df
+
+
+def load_top_chunks(top_chunks_dir: str) -> dict:
+ """
+ Загружает JSON-файлы с топ-чанками для вопросов.
+
+ Args:
+ top_chunks_dir: Директория с JSON-файлами
+
+ Returns:
+ Словарь {question_id: данные из JSON}
+ """
+ print(f"Загрузка топ-чанков из {top_chunks_dir}...")
+
+ top_chunks_data = {}
+ json_files = list(Path(top_chunks_dir).glob("question_*_top_chunks.json"))
+
+ for json_file in tqdm(json_files, desc="Загрузка JSON-файлов"):
+ try:
+ with open(json_file, 'r', encoding='utf-8') as f:
+ data = json.load(f)
+ question_id = data.get('question_id')
+ if question_id is not None:
+ top_chunks_data[question_id] = data
+ except Exception as e:
+ print(f"Ошибка при загрузке файла {json_file}: {e}")
+
+ print(f"Загружены данные для {len(top_chunks_data)} вопросов")
+
+ return top_chunks_data
+
+
+def calculate_chunk_overlap(chunk_text: str, punct_text: str) -> float:
+ """
+ Рассчитывает степень перекрытия между чанком и пунктом с использованием partial_ratio.
+
+ Args:
+ chunk_text: Текст чанка
+ punct_text: Текст пункта
+
+ Returns:
+ Коэффициент перекрытия от 0 до 1
+ """
+ # Если чанк входит в пункт, возвращаем 1.0 (полное вхождение)
+ if chunk_text in punct_text:
+ return 1.0
+
+ # Если пункт входит в чанк, возвращаем соотношение длин
+ if punct_text in chunk_text:
+ return len(punct_text) / len(chunk_text)
+
+ # Используем partial_ratio из fuzzywuzzy
+ partial_ratio_score = fuzz.partial_ratio(chunk_text, punct_text) / 100.0
+
+ return partial_ratio_score
+
+
+def find_most_similar_chunks(punct_text: str, chunks_df: pd.DataFrame, chunks_embeddings: np.ndarray, punct_embedding: np.ndarray, top_n: int = 5) -> list:
+ """
+ Находит топ-N наиболее похожих чанков для заданного пункта.
+
+ Args:
+ punct_text: Текст пункта
+ chunks_df: DataFrame с чанками
+ chunks_embeddings: Эмбеддинги чанков
+ punct_embedding: Эмбеддинг пункта
+ top_n: Количество похожих чанков (по умолчанию 5)
+
+ Returns:
+ Список словарей с информацией о похожих чанках
+ """
+ # Вычисляем косинусную близость между пунктом и всеми чанками
+ similarities = cosine_similarity([punct_embedding], chunks_embeddings)[0]
+
+ # Получаем индексы топ-N чанков по косинусной близости
+ top_indices = np.argsort(similarities)[-top_n:][::-1]
+
+ similar_chunks = []
+ for idx in top_indices:
+ chunk = chunks_df.iloc[idx]
+ overlap = calculate_chunk_overlap(chunk['text'], punct_text)
+
+ similar_chunks.append({
+ 'chunk_id': chunk['id'],
+ 'doc_name': chunk['doc_name'],
+ 'text': chunk['text'],
+ 'similarity': float(similarities[idx]),
+ 'overlap': overlap
+ })
+
+ return similar_chunks
+
+
+def analyze_missing_puncts(questions_df: pd.DataFrame, chunks_df: pd.DataFrame,
+ questions_embeddings: np.ndarray, chunks_embeddings: np.ndarray,
+ similarity_threshold: float, top_n: int = 100) -> dict:
+ """
+ Анализирует ненайденные пункты и находит для них наиболее похожие чанки.
+
+ Args:
+ questions_df: DataFrame с вопросами и пунктами
+ chunks_df: DataFrame с чанками
+ questions_embeddings: Эмбеддинги вопросов
+ chunks_embeddings: Эмбеддинги чанков
+ similarity_threshold: Порог для определения найденных пунктов
+ top_n: Количество чанков для проверки (по умолчанию 100)
+
+ Returns:
+ Словарь с результатами анализа
+ """
+ print("Анализ ненайденных пунктов...")
+
+ # Проверяем соответствие количества вопросов и эмбеддингов
+ unique_question_ids = questions_df['id'].unique()
+ if len(unique_question_ids) != questions_embeddings.shape[0]:
+ print(f"ВНИМАНИЕ: Количество уникальных ID вопросов ({len(unique_question_ids)}) не соответствует размеру массива эмбеддингов ({questions_embeddings.shape[0]}).")
+ print("Будут анализироваться только вопросы, имеющие соответствующие эмбеддинги.")
+
+ # Создаем маппинг id вопроса -> индекс в DataFrame с метаданными
+ # Используем порядковый номер в списке уникальных ID, а не порядок строк в DataFrame
+ question_id_to_idx = {qid: idx for idx, qid in enumerate(unique_question_ids)}
+
+ # Вычисляем косинусную близость между вопросами и чанками
+ similarity_matrix = cosine_similarity(questions_embeddings, chunks_embeddings)
+
+ # Результаты анализа
+ analysis_results = {}
+
+ # Обрабатываем только те вопросы, для которых у нас есть эмбеддинги
+ valid_question_ids = [qid for qid in unique_question_ids if qid in question_id_to_idx and question_id_to_idx[qid] < len(questions_embeddings)]
+
+ # Группируем датасет по id вопроса
+ for question_id in tqdm(valid_question_ids, desc="Анализ вопросов"):
+ # Получаем строки для текущего вопроса
+ question_rows = questions_df[questions_df['id'] == question_id]
+
+ # Если нет строк с таким id, пропускаем
+ if len(question_rows) == 0:
+ continue
+
+ # Получаем индекс вопроса в массиве эмбеддингов
+ question_idx = question_id_to_idx[question_id]
+
+ # Если индекс выходит за границы массива эмбеддингов, пропускаем
+ if question_idx >= questions_embeddings.shape[0]:
+ print(f"ВНИМАНИЕ: Индекс {question_idx} для вопроса {question_id} выходит за границы массива эмбеддингов размера {questions_embeddings.shape[0]}. Пропускаем.")
+ continue
+
+ # Получаем текст вопроса и пункты
+ question_text = question_rows['question'].iloc[0]
+
+ # Собираем пункты с информацией о документе
+ puncts = []
+ for _, row in question_rows.iterrows():
+ punct_doc = row.get('filename', '') if 'filename' in row else ''
+ if pd.isna(punct_doc):
+ punct_doc = ''
+ puncts.append({
+ 'text': row['text'],
+ 'doc_name': punct_doc
+ })
+
+ # Получаем связанные документы
+ relevant_docs = []
+ if 'filename' in question_rows.columns:
+ relevant_docs = [f for f in question_rows['filename'].unique() if f and not pd.isna(f)]
+ else:
+ relevant_docs = chunks_df['doc_name'].unique().tolist()
+
+ # Если для вопроса нет релевантных документов, пропускаем
+ if not relevant_docs:
+ continue
+
+ # Для отслеживания найденных и ненайденных пунктов
+ found_puncts = []
+ missing_puncts = []
+
+ # Собираем все чанки для документов вопроса
+ all_question_chunks = []
+ all_question_similarities = []
+
+ for filename in relevant_docs:
+ if not filename or pd.isna(filename):
+ continue
+
+ # Фильтруем чанки по имени файла
+ doc_chunks = chunks_df[chunks_df['doc_name'] == filename]
+
+ if doc_chunks.empty:
+ continue
+
+ # Индексы чанков для текущего файла
+ doc_chunk_indices = doc_chunks.index.tolist()
+
+ # Проверяем, что индексы чанков существуют в chunks_df
+ valid_indices = [idx for idx in doc_chunk_indices if idx in chunks_df.index]
+
+ # Получаем значения близости для чанков текущего файла
+ doc_similarities = []
+ for idx in valid_indices:
+ try:
+ chunk_loc = chunks_df.index.get_loc(idx)
+ doc_similarities.append(similarity_matrix[question_idx, chunk_loc])
+ except (KeyError, IndexError) as e:
+ print(f"Ошибка при получении индекса для чанка {idx}: {e}")
+ continue
+
+ # Добавляем чанки и их схожести к общему списку для вопроса
+ for i, idx in enumerate(valid_indices):
+ if i < len(doc_similarities): # проверяем, что у нас есть соответствующее значение similarity
+ try:
+ chunk_row = doc_chunks.loc[idx]
+ all_question_chunks.append((idx, chunk_row))
+ all_question_similarities.append(doc_similarities[i])
+ except KeyError as e:
+ print(f"Ошибка при доступе к строке с индексом {idx}: {e}")
+
+ # Если нет чанков для вопроса, пропускаем
+ if not all_question_chunks:
+ continue
+
+ # Сортируем все чанки по убыванию схожести и берем top_n
+ sorted_indices = np.argsort(all_question_similarities)[-min(top_n, len(all_question_similarities)):][::-1]
+ top_chunks = []
+ top_similarities = []
+
+ # Собираем топ-N чанков и их схожести
+ for i in sorted_indices:
+ idx, chunk = all_question_chunks[i]
+ top_chunks.append({
+ 'id': chunk['id'],
+ 'doc_name': chunk['doc_name'],
+ 'text': chunk['text']
+ })
+ top_similarities.append(all_question_similarities[i])
+
+ # Проверяем каждый пункт на наличие в топ-чанках
+ for i, punct in enumerate(puncts):
+ is_found = False
+ punct_text = punct['text']
+ punct_doc = punct['doc_name']
+
+ # Для каждого чанка из топ-N рассчитываем partial_ratio с пунктом
+ chunk_overlaps = []
+ for j, chunk in enumerate(top_chunks):
+ overlap = calculate_chunk_overlap(chunk['text'], punct_text)
+
+ # Если перекрытие больше порога, пункт найден
+ if overlap >= similarity_threshold:
+ is_found = True
+
+ # Сохраняем информацию о перекрытии для каждого чанка
+ chunk_overlaps.append({
+ 'chunk_id': chunk['id'],
+ 'doc_name': chunk['doc_name'],
+ 'text': chunk['text'],
+ 'overlap': overlap,
+ 'similarity': float(top_similarities[j])
+ })
+
+ # Если пункт найден, добавляем в список найденных
+ if is_found:
+ found_puncts.append({
+ 'index': i,
+ 'text': punct_text,
+ 'doc_name': punct_doc
+ })
+ else:
+ # Сортируем чанки по убыванию перекрытия с пунктом и берем топ-5
+ chunk_overlaps.sort(key=lambda x: x['overlap'], reverse=True)
+ top_overlaps = chunk_overlaps[:5]
+
+ missing_puncts.append({
+ 'index': i,
+ 'text': punct_text,
+ 'doc_name': punct_doc,
+ 'similar_chunks': top_overlaps
+ })
+
+ # Добавляем результаты для текущего вопроса
+ analysis_results[question_id] = {
+ 'question_id': question_id,
+ 'question_text': question_text,
+ 'found_puncts_count': len(found_puncts),
+ 'missing_puncts_count': len(missing_puncts),
+ 'total_puncts_count': len(puncts),
+ 'found_puncts': found_puncts,
+ 'missing_puncts': missing_puncts
+ }
+
+ return analysis_results
+
+
+def generate_markdown_report(analysis_results: dict, output_file: str,
+ words_per_chunk: int, overlap_words: int, model_name: str, top_n: int):
+ """
+ Генерирует отчет в формате Markdown.
+
+ Args:
+ analysis_results: Результаты анализа
+ output_file: Путь к выходному файлу
+ words_per_chunk: Размер чанка в словах
+ overlap_words: Перекрытие в словах
+ model_name: Название модели
+ top_n: Количество чанков в топе
+ """
+ print(f"Генерация отчета в формате Markdown в {output_file}...")
+
+ with open(output_file, 'w', encoding='utf-8') as f:
+ # Заголовок отчета
+ f.write(f"# Анализ ненайденных пунктов для оптимальной конфигурации чанкинга\n\n")
+
+ # Параметры анализа
+ f.write("## Параметры анализа\n\n")
+ f.write(f"- **Модель**: {model_name}\n")
+ f.write(f"- **Размер чанка**: {words_per_chunk} слов\n")
+ f.write(f"- **Перекрытие**: {overlap_words} слов ({round(overlap_words/words_per_chunk*100, 1)}%)\n")
+ f.write(f"- **Количество чанков в топе**: {top_n}\n\n")
+
+ # Сводная статистика
+ total_questions = len(analysis_results)
+ total_puncts = sum(q['total_puncts_count'] for q in analysis_results.values())
+ total_found = sum(q['found_puncts_count'] for q in analysis_results.values())
+ total_missing = sum(q['missing_puncts_count'] for q in analysis_results.values())
+
+ f.write("## Сводная статистика\n\n")
+ f.write(f"- **Всего вопросов**: {total_questions}\n")
+ f.write(f"- **Всего пунктов**: {total_puncts}\n")
+ f.write(f"- **Найдено пунктов**: {total_found} ({round(total_found/total_puncts*100, 1)}%)\n")
+ f.write(f"- **Ненайдено пунктов**: {total_missing} ({round(total_missing/total_puncts*100, 1)}%)\n\n")
+
+ # Детали по каждому вопросу
+ f.write("## Детальный анализ по вопросам\n\n")
+
+ # Сортируем вопросы по количеству ненайденных пунктов (по убыванию)
+ sorted_questions = sorted(
+ analysis_results.values(),
+ key=lambda x: x['missing_puncts_count'],
+ reverse=True
+ )
+
+ for question_data in sorted_questions:
+ question_id = question_data['question_id']
+ question_text = question_data['question_text']
+ missing_count = question_data['missing_puncts_count']
+ total_count = question_data['total_puncts_count']
+
+ # Если нет ненайденных пунктов, пропускаем
+ if missing_count == 0:
+ continue
+
+ f.write(f"### Вопрос {question_id}\n\n")
+ f.write(f"**Текст вопроса**: {question_text}\n\n")
+ f.write(f"**Статистика**: найдено {question_data['found_puncts_count']} из {total_count} пунктов ")
+ f.write(f"({round(question_data['found_puncts_count']/total_count*100, 1)}%)\n\n")
+
+ # Детали по ненайденным пунктам
+ f.write("#### Ненайденные пункты\n\n")
+
+ for i, punct in enumerate(question_data['missing_puncts']):
+ punct_text = punct['text']
+ punct_doc = punct.get('doc_name', '')
+ similar_chunks = punct['similar_chunks']
+
+ f.write(f"##### Пункт {i+1}\n\n")
+ f.write(f"**Текст пункта**: {punct_text}\n\n")
+ if punct_doc:
+ f.write(f"**Документ пункта**: {punct_doc}\n\n")
+ f.write("**Топ-5 наиболее похожих чанков**:\n\n")
+
+ # Таблица с похожими чанками
+ f.write("| № | Документ | Схожесть (с вопросом) | Перекрытие (с пунктом) | Текст чанка |\n")
+ f.write("|---|----------|----------|------------|------------|\n")
+
+ for j, chunk in enumerate(similar_chunks):
+ # Используем полный текст чанка без обрезки
+ chunk_text = chunk['text']
+
+ f.write(f"| {j+1} | {chunk['doc_name']} | {chunk['similarity']:.4f} | ")
+ f.write(f"{chunk['overlap']:.4f} | {chunk_text} |\n")
+
+ f.write("\n")
+
+ f.write("\n")
+
+ print(f"Отчет успешно сгенерирован: {output_file}")
+
+
+def main():
+ """
+ Основная функция скрипта.
+ """
+ args = parse_args()
+
+ # Загружаем датасет с вопросами
+ questions_df = load_questions_dataset(args.dataset_path)
+
+ # Загружаем чанки и эмбеддинги
+ chunks_df, chunks_embeddings, questions_embeddings, questions_meta = load_chunks_and_embeddings(
+ args.output_dir, args.words_per_chunk, args.overlap_words, args.model_name
+ )
+
+ # Анализируем ненайденные пункты
+ analysis_results = analyze_missing_puncts(
+ questions_df, chunks_df, questions_embeddings, chunks_embeddings,
+ args.similarity_threshold, args.top_n
+ )
+
+ # Генерируем отчет в формате Markdown
+ output_file = os.path.join(args.output_dir, args.markdown_file)
+ generate_markdown_report(
+ analysis_results, output_file,
+ args.words_per_chunk, args.overlap_words, args.model_name, args.top_n
+ )
+
+ print(f"Анализ ненайденных пунктов завершен. Результаты сохранены в {output_file}")
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/lib/extractor/scripts/combine_results.py b/lib/extractor/scripts/combine_results.py
new file mode 100644
index 0000000000000000000000000000000000000000..131026f36036234cdca6c4a6042937e0b414de3a
--- /dev/null
+++ b/lib/extractor/scripts/combine_results.py
@@ -0,0 +1,1352 @@
+#!/usr/bin/env python
+"""
+Скрипт для объединения результатов всех экспериментов в одну Excel-таблицу с форматированием.
+Анализирует результаты экспериментов и создает сводную таблицу с метриками в различных разрезах.
+Также строит графики через seaborn и сохраняет их в отдельную директорию.
+"""
+
+import argparse
+import glob
+import os
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
+from openpyxl import Workbook
+from openpyxl.styles import Alignment, Border, Font, PatternFill, Side
+from openpyxl.utils import get_column_letter
+from openpyxl.utils.dataframe import dataframe_to_rows
+
+
+def setup_plot_directory(plots_dir: str) -> None:
+ """
+ Создает директорию для сохранения графиков, если она не существует.
+
+ Args:
+ plots_dir: Путь к директории для графиков
+ """
+ if not os.path.exists(plots_dir):
+ os.makedirs(plots_dir)
+ print(f"Создана директория для графиков: {plots_dir}")
+ else:
+ print(f"Директория для графиков: {plots_dir}")
+
+
+def parse_args():
+ """Парсит аргументы командной строки."""
+ parser = argparse.ArgumentParser(description="Объединение результатов экспериментов в одну Excel-таблицу")
+
+ parser.add_argument("--results-dir", type=str, default="data",
+ help="Директория с результатами экспериментов (по умолчанию: data)")
+ parser.add_argument("--output-file", type=str, default="combined_results.xlsx",
+ help="Путь к выходному Excel-файлу (по умолчанию: combined_results.xlsx)")
+ parser.add_argument("--plots-dir", type=str, default="plots",
+ help="Директория для сохранения графиков (по умолчанию: plots)")
+
+ return parser.parse_args()
+
+
+def parse_file_name(file_name: str) -> dict:
+ """
+ Парсит имя файла и извлекает параметры эксперимента.
+
+ Args:
+ file_name: Имя файла для парсинга
+
+ Returns:
+ Словарь с параметрами (words_per_chunk, overlap_words, model) или None при ошибке
+ """
+ try:
+ # Извлекаем параметры из имени файла
+ parts = file_name.split('_')
+ if len(parts) < 4:
+ return None
+
+ # Ищем части с w (words) и o (overlap)
+ words_part = None
+ overlap_part = None
+
+ for part in parts:
+ if part.startswith('w') and part[1:].isdigit():
+ words_part = part[1:]
+ elif part.startswith('o') and part[1:].isdigit():
+ # Убираем потенциальную часть .csv или .xlsx из overlap_part
+ overlap_part = part[1:].split('.')[0]
+
+ if words_part is None or overlap_part is None:
+ return None
+
+ # Пытаемся извлечь имя модели из оставшейся части имени файла
+ model_part = file_name.split(f"_w{words_part}_o{overlap_part}_", 1)
+ if len(model_part) < 2:
+ return None
+
+ # Получаем имя модели и удаляем возможное расширение файла
+ model_name_parts = model_part[1].split('.')
+ if len(model_name_parts) > 1:
+ model_name_parts = model_name_parts[:-1]
+
+ model_name_parts = '_'.join(model_name_parts).split('_')
+ model_name = '/'.join(model_name_parts)
+
+ return {
+ 'words_per_chunk': int(words_part),
+ 'overlap_words': int(overlap_part),
+ 'model': model_name,
+ 'overlap_percentage': round(int(overlap_part) / int(words_part) * 100, 1)
+ }
+ except Exception as e:
+ print(f"Ошибка при парсинге файла {file_name}: {e}")
+ return None
+
+
+def load_data_files(results_dir: str, pattern: str, file_type: str, load_function) -> pd.DataFrame:
+ """
+ Общая функция для загрузки файлов данных с определенным паттерном имени.
+
+ Args:
+ results_dir: Директория с результатами
+ pattern: Glob-паттерн для поиска файлов
+ file_type: Тип файлов для сообщений (напр. "результатов", "метрик")
+ load_function: Функция для загрузки конкретного типа файла
+
+ Returns:
+ DataFrame с объединенными данными или None при ошибке
+ """
+ print(f"Загрузка {file_type} из {results_dir}...")
+
+ # Ищем все файлы с указанным паттерном
+ data_files = glob.glob(os.path.join(results_dir, pattern))
+
+ if not data_files:
+ print(f"В директории {results_dir} не найдены файлы {file_type}")
+ return None
+
+ print(f"Найдено {len(data_files)} файлов {file_type}")
+
+ all_data = []
+
+ for file_path in data_files:
+ # Извлекаем информацию о стратегии и модели из имени файла
+ file_name = os.path.basename(file_path)
+ print(f"Обрабатываю файл: {file_name}")
+
+ # Парсим параметры из имени файла
+ params = parse_file_name(file_name)
+
+ if params is None:
+ print(f"Пропуск файла {file_name}: не удалось извлечь параметры")
+ continue
+
+ words_part = params['words_per_chunk']
+ overlap_part = params['overlap_words']
+ model_name = params['model']
+ overlap_percentage = params['overlap_percentage']
+
+ print(f" Параметры: words={words_part}, overlap={overlap_part}, model={model_name}")
+
+ try:
+ # Загружаем данные, используя переданную функцию
+ df = load_function(file_path)
+
+ # Добавляем информацию о стратегии и модели
+ df['model'] = model_name
+ df['words_per_chunk'] = words_part
+ df['overlap_words'] = overlap_part
+ df['overlap_percentage'] = overlap_percentage
+
+ all_data.append(df)
+ except Exception as e:
+ print(f"Ошибка при обработке файла {file_path}: {e}")
+
+ if not all_data:
+ print(f"Не удалось загрузить ни один файл {file_type}")
+ return None
+
+ # Объединяем все данные
+ combined_data = pd.concat(all_data, ignore_index=True)
+
+ return combined_data
+
+
+def load_results_files(results_dir: str) -> pd.DataFrame:
+ """
+ Загружает все файлы результатов из указанной директории.
+
+ Args:
+ results_dir: Директория с результатами
+
+ Returns:
+ DataFrame с объединенными результатами
+ """
+ # Используем общую функцию для загрузки CSV файлов
+ data = load_data_files(
+ results_dir,
+ "results_*.csv",
+ "результатов",
+ lambda f: pd.read_csv(f)
+ )
+
+ if data is None:
+ raise ValueError("Не удалось загрузить файлы с результатами")
+
+ return data
+
+
+def load_question_metrics_files(results_dir: str) -> pd.DataFrame:
+ """
+ Загружает все файлы с метриками по вопросам из указанной директории.
+
+ Args:
+ results_dir: Директория с результатами
+
+ Returns:
+ DataFrame с объединенными метриками по вопросам или None, если файлов нет
+ """
+ # Используем общую функцию для загрузки Excel файлов
+ return load_data_files(
+ results_dir,
+ "question_metrics_*.xlsx",
+ "метрик по вопросам",
+ lambda f: pd.read_excel(f)
+ )
+
+
+def prepare_summary_by_model_top_n(df: pd.DataFrame, macro_metrics: pd.DataFrame = None) -> pd.DataFrame:
+ """
+ Подготавливает сводную таблицу по моделям и top_n значениям.
+ Если доступны macro метрики, они также включаются в сводную таблицу.
+
+ Args:
+ df: DataFrame с объединенными результатами
+ macro_metrics: DataFrame с macro метриками (опционально)
+
+ Returns:
+ DataFrame со сводной таблицей
+ """
+ # Определяем группировочные колонки и метрики
+ group_by_columns = ['model', 'top_n']
+ metrics = ['text_precision', 'text_recall', 'text_f1', 'doc_precision', 'doc_recall', 'doc_f1']
+
+ # Используем общую функцию для подготовки сводки
+ return prepare_summary(df, group_by_columns, metrics, macro_metrics)
+
+
+def prepare_summary_by_chunking_params_top_n(df: pd.DataFrame, macro_metrics: pd.DataFrame = None) -> pd.DataFrame:
+ """
+ Подготавливает сводную таблицу по параметрам чанкинга и top_n значениям.
+ Если доступны macro метрики, они также включаются в сводную таблицу.
+
+ Args:
+ df: DataFrame с объединенными результатами
+ macro_metrics: DataFrame с macro метриками (опционально)
+
+ Returns:
+ DataFrame со сводной таблицей
+ """
+ # Определяем группировочные колонки и метрики
+ group_by_columns = ['words_per_chunk', 'overlap_words', 'top_n']
+ metrics = ['text_precision', 'text_recall', 'text_f1', 'doc_precision', 'doc_recall', 'doc_f1']
+
+ # Используем общую функцию для подготовки сводки
+ return prepare_summary(df, group_by_columns, metrics, macro_metrics)
+
+
+def prepare_summary(df: pd.DataFrame, group_by_columns: list, metrics: list, macro_metrics: pd.DataFrame = None) -> pd.DataFrame:
+ """
+ Общая функция для подготовки сводной таблицы по указанным группировочным колонкам.
+ Если доступны macro метрики, они также включаются в сводную таблицу.
+
+ Args:
+ df: DataFrame с объединенными результатами
+ group_by_columns: Колонки для группировки
+ metrics: Список метрик для расчета среднего
+ macro_metrics: DataFrame с macro метриками (опционально)
+
+ Returns:
+ DataFrame со сводной таблицей
+ """
+ # Группируем по указанным колонкам, вычисляем средние значения метрик
+ summary = df.groupby(group_by_columns).agg({
+ metric: 'mean' for metric in metrics
+ }).reset_index()
+
+ # Если среди группировочных колонок есть 'overlap_words' и 'words_per_chunk',
+ # добавляем процент перекрытия
+ if 'overlap_words' in group_by_columns and 'words_per_chunk' in group_by_columns:
+ summary['overlap_percentage'] = (summary['overlap_words'] / summary['words_per_chunk'] * 100).round(1)
+
+ # Если доступны macro метрики, объединяем их с summary
+ if macro_metrics is not None:
+ # Преобразуем метрики в macro_метрики
+ macro_metric_names = [f"macro_{metric}" for metric in metrics]
+
+ # Группируем macro метрики по тем же колонкам
+ macro_summary = macro_metrics.groupby(group_by_columns).agg({
+ metric: 'mean' for metric in macro_metric_names
+ }).reset_index()
+
+ # Если нужно, добавляем процент перекрытия для согласованности
+ if 'overlap_words' in group_by_columns and 'words_per_chunk' in group_by_columns:
+ macro_summary['overlap_percentage'] = (macro_summary['overlap_words'] / macro_summary['words_per_chunk'] * 100).round(1)
+ merge_on = group_by_columns + ['overlap_percentage']
+ else:
+ merge_on = group_by_columns
+
+ # Объединяем с основной сводкой
+ summary = pd.merge(summary, macro_summary, on=merge_on, how='left')
+
+ # Сортируем по группировочным колонкам
+ summary = summary.sort_values(group_by_columns)
+
+ # Округляем метрики до 4 знаков после запятой
+ for col in summary.columns:
+ if any(col.endswith(suffix) for suffix in ['precision', 'recall', 'f1']):
+ summary[col] = summary[col].round(4)
+
+ return summary
+
+
+def prepare_best_configurations(df: pd.DataFrame, macro_metrics: pd.DataFrame = None) -> pd.DataFrame:
+ """
+ Подготавливает таблицу с лучшими конфигурациями для каждой модели и различных top_n.
+ Выбирает конфигурацию только на основе macro_text_recall и text_recall (weighted),
+ игнорируя F1 метрики как менее важные.
+
+ Args:
+ df: DataFrame с объединенными результатами
+ macro_metrics: DataFrame с macro метриками (опционально)
+
+ Returns:
+ DataFrame с лучшими конфигурациями
+ """
+ # Выбираем ключевые значения top_n
+ key_top_n = [10, 20, 50, 100]
+
+ # Определяем источник метрик и акцентируем только на recall-метриках
+ if macro_metrics is not None:
+ print("Выбор лучших конфигураций на основе macro метрик (macro_text_recall)")
+ metrics_source = macro_metrics
+ text_recall_metric = 'macro_text_recall'
+ doc_recall_metric = 'macro_doc_recall'
+ else:
+ print("Выбор лучших конфигураций на основе weighted метрик (text_recall)")
+ metrics_source = df
+ text_recall_metric = 'text_recall'
+ doc_recall_metric = 'doc_recall'
+
+ # Фильтруем только по ключевым значениям top_n
+ filtered_df = metrics_source[metrics_source['top_n'].isin(key_top_n)]
+
+ # Для каждой модели и top_n находим конфигурацию только с лучшим recall
+ best_configs = []
+
+ for model in metrics_source['model'].unique():
+ for top_n in key_top_n:
+ model_top_n_df = filtered_df[(filtered_df['model'] == model) & (filtered_df['top_n'] == top_n)]
+
+ if len(model_top_n_df) == 0:
+ continue
+
+ # Находим конфигурацию с лучшим text_recall
+ best_text_recall_idx = model_top_n_df[text_recall_metric].idxmax()
+ best_text_recall_config = model_top_n_df.loc[best_text_recall_idx].copy()
+ best_text_recall_config['metric_type'] = 'text_recall'
+
+ # Находим конфигурацию с лучшим doc_recall
+ best_doc_recall_idx = model_top_n_df[doc_recall_metric].idxmax()
+ best_doc_recall_config = model_top_n_df.loc[best_doc_recall_idx].copy()
+ best_doc_recall_config['metric_type'] = 'doc_recall'
+
+ best_configs.append(best_text_recall_config)
+ best_configs.append(best_doc_recall_config)
+
+ if not best_configs:
+ return pd.DataFrame()
+
+ best_configs_df = pd.DataFrame(best_configs)
+
+ # Выбираем и сортируем нужные столбцы
+ cols_to_keep = ['model', 'top_n', 'metric_type', 'words_per_chunk', 'overlap_words', 'overlap_percentage']
+
+ # Добавляем столбцы метрик в зависимости от того, какие доступны
+ if macro_metrics is not None:
+ # Для macro метрик сначала выбираем recall-метрики
+ recall_cols = [col for col in best_configs_df.columns if col.endswith('recall')]
+ # Затем добавляем остальные метрики
+ other_cols = [col for col in best_configs_df.columns if any(col.endswith(m) for m in
+ ['precision', 'f1']) and col.startswith('macro_')]
+ metric_cols = recall_cols + other_cols
+ else:
+ # Для weighted метрик сначала выбираем recall-метрики
+ recall_cols = [col for col in best_configs_df.columns if col.endswith('recall')]
+ # Затем добавляем остальные метрики
+ other_cols = [col for col in best_configs_df.columns if any(col.endswith(m) for m in
+ ['precision', 'f1']) and not col.startswith('macro_')]
+ metric_cols = recall_cols + other_cols
+
+ result = best_configs_df[cols_to_keep + metric_cols].sort_values(['model', 'top_n', 'metric_type'])
+
+ return result
+
+
+def get_grouping_columns(sheet) -> dict:
+ """
+ Определяет подходящие колонки для группировки данных на листе.
+
+ Args:
+ sheet: Лист Excel
+
+ Returns:
+ Словарь с данными о группировке или None
+ """
+ # Возможные варианты группировки
+ grouping_possibilities = [
+ {'columns': ['model', 'words_per_chunk', 'overlap_words']},
+ {'columns': ['model']},
+ {'columns': ['words_per_chunk', 'overlap_words']},
+ {'columns': ['top_n']},
+ {'columns': ['model', 'top_n', 'metric_type']}
+ ]
+
+ # Для каждого варианта группировки проверяем наличие всех колонок
+ for grouping in grouping_possibilities:
+ column_indices = {}
+ all_columns_present = True
+
+ for column_name in grouping['columns']:
+ column_idx = None
+ for col_idx, cell in enumerate(sheet[1], start=1):
+ if cell.value == column_name:
+ column_idx = col_idx
+ break
+
+ if column_idx is None:
+ all_columns_present = False
+ break
+ else:
+ column_indices[column_name] = column_idx
+
+ if all_columns_present:
+ return {
+ 'columns': grouping['columns'],
+ 'indices': column_indices
+ }
+
+ return None
+
+
+def apply_header_formatting(sheet):
+ """
+ Применяет форматирование к заголовкам.
+
+ Args:
+ sheet: Лист Excel
+ """
+ # Форматирование заголовков
+ for cell in sheet[1]:
+ cell.font = Font(bold=True)
+ cell.fill = PatternFill(start_color="D9D9D9", end_color="D9D9D9", fill_type="solid")
+ cell.alignment = Alignment(horizontal='center', vertical='center', wrap_text=True)
+
+
+def adjust_column_width(sheet):
+ """
+ Настраивает ширину столбцов на основе содержимого.
+
+ Args:
+ sheet: Лист Excel
+ """
+ # Авторазмер столбцов
+ for column in sheet.columns:
+ max_length = 0
+ column_letter = get_column_letter(column[0].column)
+
+ for cell in column:
+ if cell.value:
+ try:
+ if len(str(cell.value)) > max_length:
+ max_length = len(str(cell.value))
+ except:
+ pass
+
+ adjusted_width = (max_length + 2) * 1.1
+ sheet.column_dimensions[column_letter].width = adjusted_width
+
+
+def apply_cell_formatting(sheet):
+ """
+ Применяет форматирование к ячейкам (границы, выравнивание и т.д.).
+
+ Args:
+ sheet: Лист Excel
+ """
+ # Тонкие границы для всех ячеек
+ thin_border = Border(
+ left=Side(style='thin'),
+ right=Side(style='thin'),
+ top=Side(style='thin'),
+ bottom=Side(style='thin')
+ )
+
+ for row in sheet.iter_rows(min_row=1, max_row=sheet.max_row, min_col=1, max_col=sheet.max_column):
+ for cell in row:
+ cell.border = thin_border
+
+ # Форматирование числовых значений
+ numeric_columns = [
+ 'text_precision', 'text_recall', 'text_f1',
+ 'doc_precision', 'doc_recall', 'doc_f1',
+ 'macro_text_precision', 'macro_text_recall', 'macro_text_f1',
+ 'macro_doc_precision', 'macro_doc_recall', 'macro_doc_f1'
+ ]
+
+ for col_idx, header in enumerate(sheet[1], start=1):
+ if header.value in numeric_columns or (header.value and str(header.value).endswith(('precision', 'recall', 'f1'))):
+ for row_idx in range(2, sheet.max_row + 1):
+ cell = sheet.cell(row=row_idx, column=col_idx)
+ if isinstance(cell.value, (int, float)):
+ cell.number_format = '0.0000'
+
+ # Выравнивание для всех ячеек
+ for row in sheet.iter_rows(min_row=2, max_row=sheet.max_row, min_col=1, max_col=sheet.max_column):
+ for cell in row:
+ cell.alignment = Alignment(horizontal='center', vertical='center')
+
+
+def apply_group_formatting(sheet, grouping):
+ """
+ Применяет форматирование к группам строк.
+
+ Args:
+ sheet: Лист Excel
+ grouping: Словарь с данными о группировке
+ """
+ if not grouping or sheet.max_row <= 1:
+ return
+
+ # Для каждой строки проверяем изменение значений группировочных колонок
+ last_values = {column: None for column in grouping['columns']}
+
+ # Применяем жирную верхнюю границу к первой строке данных
+ for col_idx in range(1, sheet.max_column + 1):
+ cell = sheet.cell(row=2, column=col_idx)
+ cell.border = Border(
+ left=cell.border.left,
+ right=cell.border.right,
+ top=Side(style='thick'),
+ bottom=cell.border.bottom
+ )
+
+ for row_idx in range(2, sheet.max_row + 1):
+ current_values = {}
+ for column in grouping['columns']:
+ col_idx = grouping['indices'][column]
+ current_values[column] = sheet.cell(row=row_idx, column=col_idx).value
+
+ # Если значения изменились, добавляем жирные границы
+ values_changed = False
+ for column in grouping['columns']:
+ if current_values[column] != last_values[column]:
+ values_changed = True
+ break
+
+ if values_changed and row_idx > 2:
+ # Жирная верхняя граница для текущей строки
+ for col_idx in range(1, sheet.max_column + 1):
+ cell = sheet.cell(row=row_idx, column=col_idx)
+ cell.border = Border(
+ left=cell.border.left,
+ right=cell.border.right,
+ top=Side(style='thick'),
+ bottom=cell.border.bottom
+ )
+
+ # Жирная нижняя граница для предыдущей строки
+ for col_idx in range(1, sheet.max_column + 1):
+ cell = sheet.cell(row=row_idx-1, column=col_idx)
+ cell.border = Border(
+ left=cell.border.left,
+ right=cell.border.right,
+ top=cell.border.top,
+ bottom=Side(style='thick')
+ )
+
+ # Запоминаем текущие значения для следующей итерации
+ for column in grouping['columns']:
+ last_values[column] = current_values[column]
+
+ # Добавляем жирную нижнюю границу для последней строки
+ for col_idx in range(1, sheet.max_column + 1):
+ cell = sheet.cell(row=sheet.max_row, column=col_idx)
+ cell.border = Border(
+ left=cell.border.left,
+ right=cell.border.right,
+ top=cell.border.top,
+ bottom=Side(style='thick')
+ )
+
+
+def apply_formatting(workbook: Workbook) -> None:
+ """
+ Применяет форматирование к Excel-файлу.
+ Добавляет автофильтры для всех столбцов и улучшает визуальное представление.
+
+ Args:
+ workbook: Workbook-объект openpyxl
+ """
+ for sheet_name in workbook.sheetnames:
+ sheet = workbook[sheet_name]
+
+ # Добавляем автофильтры для всех столбцов
+ if sheet.max_row > 1: # Проверяем, что в листе есть данные
+ sheet.auto_filter.ref = sheet.dimensions
+
+ # Применяем форматирование
+ apply_header_formatting(sheet)
+ adjust_column_width(sheet)
+ apply_cell_formatting(sheet)
+
+ # Определяем группирующие колонки и применяем форматирование к группам
+ grouping = get_grouping_columns(sheet)
+ if grouping:
+ apply_group_formatting(sheet, grouping)
+
+
+def create_model_comparison_plot(df: pd.DataFrame, metrics: list | str, top_n: int, plots_dir: str) -> None:
+ """
+ Создает график сравнения моделей по указанным метрикам для заданного top_n.
+
+ Args:
+ df: DataFrame с данными
+ metrics: Список метрик или одна метрика для сравнения
+ top_n: Значение top_n для фильтрации
+ plots_dir: Директория для сохранения графиков
+ """
+ if isinstance(metrics, str):
+ metrics = [metrics]
+
+ # Фильтруем данные
+ filtered_df = df[df['top_n'] == top_n]
+
+ if len(filtered_df) == 0:
+ print(f"Нет данных для top_n={top_n}")
+ return
+
+ # Определяем тип метрик (macro или weighted)
+ metrics_type = "macro" if metrics[0].startswith("macro_") else "weighted"
+
+ # Создаем фигуру с несколькими подграфиками
+ fig, axes = plt.subplots(1, len(metrics), figsize=(6 * len(metrics), 8))
+
+ # Если только одна метрика, преобразуем axes в список для единообразного обращения
+ if len(metrics) == 1:
+ axes = [axes]
+
+ # Для каждой метрики создаем subplot
+ for i, metric in enumerate(metrics):
+ # Группируем данные по модели
+ columns_to_agg = {metric: 'mean'}
+ model_data = filtered_df.groupby('model').agg(columns_to_agg).reset_index()
+
+ # Сортируем по значению метрики (по убыванию)
+ model_data = model_data.sort_values(metric, ascending=False)
+
+ # Определяем цветовую схему
+ palette = sns.color_palette("viridis", len(model_data))
+
+ # Строим столбчатую диаграмму на соответствующем subplot
+ ax = sns.barplot(x='model', y=metric, data=model_data, palette=palette, ax=axes[i])
+
+ # Добавляем значения над столбцами
+ for j, v in enumerate(model_data[metric]):
+ ax.text(j, v + 0.01, f"{v:.4f}", ha='center', fontsize=8)
+
+ # Устанавливаем заголовок и метки осей
+ ax.set_title(f"{metric} (top_n={top_n})", fontsize=12)
+ ax.set_xlabel("Модель", fontsize=10)
+ ax.set_ylabel(f"{metric}", fontsize=10)
+
+ # Поворачиваем подписи по оси X для лучшей читаемости
+ ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha='right', fontsize=8)
+
+ # Настраиваем макет
+ plt.tight_layout()
+
+ # Сохраняем график
+ metric_names = '_'.join([m.replace('macro_', '') for m in metrics])
+ file_name = f"model_comparison_{metrics_type}_{metric_names}_top{top_n}.png"
+ plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
+ plt.close()
+
+ print(f"Создан график сравнения моделей: {file_name}")
+
+
+def create_top_n_plot(df: pd.DataFrame, models: list | str, metric: str, plots_dir: str) -> None:
+ """
+ Создает график зависимости метрики от top_n для заданных моделей.
+
+ Args:
+ df: DataFrame с данными
+ models: Список моделей или одна модель для сравнения
+ metric: Название метрики
+ plots_dir: Директория для сохранения графиков
+ """
+ if isinstance(models, str):
+ models = [models]
+
+ # Создаем фигуру
+ plt.figure(figsize=(12, 8))
+
+ # Определяем цветовую схему
+ palette = sns.color_palette("viridis", len(models))
+
+ # Ограничиваем количество моделей для читаемости
+ if len(models) > 5:
+ models = models[:5]
+ print("Слишком много моделей для графика, ограничиваем до 5")
+
+ # Для каждой модели строим линию
+ for i, model in enumerate(models):
+ # Находим наиболее часто используемые параметры чанкинга для этой модели
+ model_df = df[df['model'] == model]
+
+ if len(model_df) == 0:
+ print(f"Нет данных для модели {model}")
+ continue
+
+ # Группируем по параметрам чанкинга и подсчитываем частоту
+ common_configs = model_df.groupby(['words_per_chunk', 'overlap_words']).size().reset_index(name='count')
+
+ if len(common_configs) == 0:
+ continue
+
+ # Берем наиболее частую конфигурацию
+ common_config = common_configs.sort_values('count', ascending=False).iloc[0]
+
+ # Фильтруем для этой конфигурации
+ config_df = model_df[
+ (model_df['words_per_chunk'] == common_config['words_per_chunk']) &
+ (model_df['overlap_words'] == common_config['overlap_words'])
+ ].sort_values('top_n')
+
+ if len(config_df) <= 1:
+ continue
+
+ # Строим линию
+ plt.plot(config_df['top_n'], config_df[metric], marker='o', linewidth=2,
+ label=f"{model} (w={common_config['words_per_chunk']}, o={common_config['overlap_words']})",
+ color=palette[i])
+
+ # Добавляем легенду, заголовок и метки осей
+ plt.legend(title="Модель (параметры)", fontsize=10, loc='best')
+ plt.title(f"Зависимость {metric} от top_n для разных моделей", fontsize=16)
+ plt.xlabel("top_n", fontsize=14)
+ plt.ylabel(metric, fontsize=14)
+
+ # Включаем сетку
+ plt.grid(True, linestyle='--', alpha=0.7)
+
+ # Настраиваем макет
+ plt.tight_layout()
+
+ # Сохраняем график
+ is_macro = "macro" if "macro" in metric else "weighted"
+ file_name = f"top_n_comparison_{is_macro}_{metric.replace('macro_', '')}.png"
+ plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
+ plt.close()
+
+ print(f"Создан график зависимости от top_n: {file_name}")
+
+
+def create_chunk_size_plot(df: pd.DataFrame, model: str, metrics: list | str, top_n: int, plots_dir: str) -> None:
+ """
+ Создает график зависимости метрик от размера чанка для заданной модели и top_n.
+
+ Args:
+ df: DataFrame с данными
+ model: Название модели
+ metrics: Список метрик или одна метрика
+ top_n: Значение top_n
+ plots_dir: Директория для сохранения графиков
+ """
+ if isinstance(metrics, str):
+ metrics = [metrics]
+
+ # Фильтруем данные
+ filtered_df = df[(df['model'] == model) & (df['top_n'] == top_n)]
+
+ if len(filtered_df) <= 1:
+ print(f"Недостаточно данных для модели {model} и top_n={top_n}")
+ return
+
+ # Создаем фигуру
+ plt.figure(figsize=(14, 8))
+
+ # Определяем цветовую схему для метрик
+ palette = sns.color_palette("viridis", len(metrics))
+
+ # Группируем по размеру чанка и проценту перекрытия
+ # Вычисляем среднее только для указанных метрик, а не для всех столбцов
+ columns_to_agg = {metric: 'mean' for metric in metrics}
+ chunk_data = filtered_df.groupby(['words_per_chunk', 'overlap_percentage']).agg(columns_to_agg).reset_index()
+
+ # Получаем уникальные значения процента перекрытия
+ overlap_percentages = sorted(chunk_data['overlap_percentage'].unique())
+
+ # Настраиваем маркеры и линии для разных перекрытий
+ markers = ['o', 's', '^', 'D', 'x', '*']
+
+ # Для каждого перекрытия строим линии с разными метриками
+ for i, overlap in enumerate(overlap_percentages):
+ subset = chunk_data[chunk_data['overlap_percentage'] == overlap].sort_values('words_per_chunk')
+
+ for j, metric in enumerate(metrics):
+ plt.plot(subset['words_per_chunk'], subset[metric],
+ marker=markers[i % len(markers)], linewidth=2,
+ label=f"{metric}, overlap={overlap}%",
+ color=palette[j])
+
+ # Добавляем легенду и заголовок
+ plt.legend(title="Метрика и перекрытие", fontsize=10, loc='best')
+ plt.title(f"Зависимость метрик от размера чанка для {model} (top_n={top_n})", fontsize=16)
+ plt.xlabel("Размер чанка (слов)", fontsize=14)
+ plt.ylabel("Значение метрики", fontsize=14)
+
+ # Включаем сетку
+ plt.grid(True, linestyle='--', alpha=0.7)
+
+ # Настраиваем макет
+ plt.tight_layout()
+
+ # Сохраняем график
+ metrics_type = "macro" if metrics[0].startswith("macro_") else "weighted"
+ model_name = model.replace('/', '_')
+ metric_names = '_'.join([m.replace('macro_', '') for m in metrics])
+ file_name = f"chunk_size_{metrics_type}_{metric_names}_{model_name}_top{top_n}.png"
+ plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
+ plt.close()
+
+ print(f"Создан график зависимости от размера чанка: {file_name}")
+
+
+def create_heatmap(df: pd.DataFrame, models: list | str, metric: str, top_n: int, plots_dir: str) -> None:
+ """
+ Создает тепловые карты зависимости метрики от размера чанка и процента перекрытия
+ для заданных моделей.
+
+ Args:
+ df: DataFrame с данными
+ models: Список моделей или одна модель
+ metric: Название метрики
+ top_n: Значение top_n
+ plots_dir: Директория для сохранения графиков
+ """
+ if isinstance(models, str):
+ models = [models]
+
+ # Ограничиваем количество моделей для наглядности
+ if len(models) > 4:
+ models = models[:4]
+
+ # Создаем фигуру с подграфиками
+ fig, axes = plt.subplots(1, len(models), figsize=(6 * len(models), 6), squeeze=False)
+
+ # Для каждой модели создаем тепловую карту
+ for i, model in enumerate(models):
+ # Фильтруем данные для указанной модели и top_n
+ filtered_df = df[(df['model'] == model) & (df['top_n'] == top_n)]
+
+ # Проверяем, достаточно ли данных для построения тепловой карты
+ chunk_sizes = filtered_df['words_per_chunk'].unique()
+ overlap_percentages = filtered_df['overlap_percentage'].unique()
+
+ if len(chunk_sizes) <= 1 or len(overlap_percentages) <= 1:
+ print(f"Недостаточно данных для построения тепловой карты для модели {model} и top_n={top_n}")
+ # Пропускаем этот subplot
+ axes[0, i].text(0.5, 0.5, f"Недостаточно данных для {model}",
+ horizontalalignment='center', verticalalignment='center')
+ axes[0, i].set_title(model)
+ axes[0, i].axis('off')
+ continue
+
+ # Создаем сводную таблицу для тепловой карты, используя только нужную метрику
+ # Сначала выберем только колонки для pivot_table
+ pivot_columns = ['words_per_chunk', 'overlap_percentage', metric]
+ pivot_df = filtered_df[pivot_columns].copy()
+
+ # Теперь создаем сводную таблицу
+ pivot_data = pivot_df.pivot_table(
+ index='words_per_chunk',
+ columns='overlap_percentage',
+ values=metric,
+ aggfunc='mean'
+ )
+
+ # Строим тепловую карту
+ sns.heatmap(pivot_data, annot=True, fmt=".4f", cmap="viridis",
+ linewidths=.5, annot_kws={"size": 8}, ax=axes[0, i])
+
+ # Устанавливаем заголовок и метки осей
+ axes[0, i].set_title(model, fontsize=12)
+ axes[0, i].set_xlabel("Процент перекрытия (%)", fontsize=10)
+ axes[0, i].set_ylabel("Размер чанка (слов)", fontsize=10)
+
+ # Добавляем общий заголовок
+ plt.suptitle(f"Тепловые карты {metric} для разных моделей (top_n={top_n})", fontsize=16)
+
+ # Настраиваем макет
+ plt.tight_layout(rect=[0, 0, 1, 0.96]) # Оставляем место для общего заголовка
+
+ # Сохраняем график
+ is_macro = "macro" if "macro" in metric else "weighted"
+ file_name = f"heatmap_{is_macro}_{metric.replace('macro_', '')}_top{top_n}.png"
+ plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
+ plt.close()
+
+ print(f"Созданы тепловые карты: {file_name}")
+
+
+def find_best_combinations(df: pd.DataFrame, metrics: list | str = None) -> pd.DataFrame:
+ """
+ Находит наилучшие комбинации параметров на основе агрегированных recall-метрик.
+
+ Args:
+ df: DataFrame с данными
+ metrics: Список метрик для анализа или None (тогда используются все recall-метрики)
+
+ Returns:
+ DataFrame с лучшими комбинациями параметров
+ """
+ if metrics is None:
+ # По умолчанию выбираем все метрики с "recall" в названии
+ metrics = [col for col in df.columns if "recall" in col]
+ elif isinstance(metrics, str):
+ metrics = [metrics]
+
+ print(f"Поиск лучших комбинаций на основе метрик: {metrics}")
+
+ # Создаем новую метрику - сумму всех указанных recall-метрик
+ df_copy = df.copy()
+ df_copy['combined_recall'] = df_copy[metrics].sum(axis=1)
+
+ # Находим лучшие комбинации для различных значений top_n
+ best_combinations = []
+
+ for top_n in df_copy['top_n'].unique():
+ top_n_df = df_copy[df_copy['top_n'] == top_n]
+
+ if len(top_n_df) == 0:
+ continue
+
+ # Находим строку с максимальным combined_recall
+ best_idx = top_n_df['combined_recall'].idxmax()
+ best_row = top_n_df.loc[best_idx].copy()
+ best_row['best_for_top_n'] = top_n
+
+ best_combinations.append(best_row)
+
+ # Находим лучшие комбинации для разных моделей
+ for model in df_copy['model'].unique():
+ model_df = df_copy[df_copy['model'] == model]
+
+ if len(model_df) == 0:
+ continue
+
+ # Находим строку с максимальным combined_recall
+ best_idx = model_df['combined_recall'].idxmax()
+ best_row = model_df.loc[best_idx].copy()
+ best_row['best_for_model'] = model
+
+ best_combinations.append(best_row)
+
+ # Находим лучшие комбинации для разных размеров чанков
+ for chunk_size in df_copy['words_per_chunk'].unique():
+ chunk_df = df_copy[df_copy['words_per_chunk'] == chunk_size]
+
+ if len(chunk_df) == 0:
+ continue
+
+ # Находим строку с максимальным combined_recall
+ best_idx = chunk_df['combined_recall'].idxmax()
+ best_row = chunk_df.loc[best_idx].copy()
+ best_row['best_for_chunk_size'] = chunk_size
+
+ best_combinations.append(best_row)
+
+ # Находим абсолютно лучшую комбинацию
+ if len(df_copy) > 0:
+ best_idx = df_copy['combined_recall'].idxmax()
+ best_row = df_copy.loc[best_idx].copy()
+ best_row['absolute_best'] = True
+
+ best_combinations.append(best_row)
+
+ if not best_combinations:
+ return pd.DataFrame()
+
+ result = pd.DataFrame(best_combinations)
+
+ # Сортируем по combined_recall (по убыванию)
+ result = result.sort_values('combined_recall', ascending=False)
+
+ print(f"Найдено {len(result)} лучших комбинаций")
+
+ return result
+
+
+def create_best_combinations_plot(best_df: pd.DataFrame, metrics: list | str, plots_dir: str) -> None:
+ """
+ Создает график сравнения лучших комбинаций параметров.
+
+ Args:
+ best_df: DataFrame с лучшими комбинациями
+ metrics: Список метрик для визуализации
+ plots_dir: Директория для сохранения графиков
+ """
+ if isinstance(metrics, str):
+ metrics = [metrics]
+
+ if len(best_df) == 0:
+ print("Нет данных для построения графика лучших комбинаций")
+ return
+
+ # Создаем новый признак для идентификации комбинаций
+ best_df['combo_label'] = best_df.apply(
+ lambda row: f"{row['model']} (w={row['words_per_chunk']}, o={row['overlap_words']}, top_n={row['top_n']})",
+ axis=1
+ )
+
+ # Берем только лучшие N комбинаций для читаемости
+ max_combos = 10
+ if len(best_df) > max_combos:
+ plot_df = best_df.head(max_combos).copy()
+ print(f"Ограничиваем график до {max_combos} лучших комбинаций")
+ else:
+ plot_df = best_df.copy()
+
+ # Создаем длинный формат данных для seaborn
+ plot_data = plot_df.melt(
+ id_vars=['combo_label', 'combined_recall'],
+ value_vars=metrics,
+ var_name='metric',
+ value_name='value'
+ )
+
+ # Сортируем по суммарному recall (комбинации) и метрике (для группировки)
+ plot_data = plot_data.sort_values(['combined_recall', 'metric'], ascending=[False, True])
+
+ # Создаем фигуру для графика
+ plt.figure(figsize=(14, 10))
+
+ # Создаем bar plot
+ sns.barplot(
+ x='combo_label',
+ y='value',
+ hue='metric',
+ data=plot_data,
+ palette='viridis'
+ )
+
+ # Настраиваем оси и заголовок
+ plt.title('Лучшие комбинации параметров по recall-метрикам', fontsize=16)
+ plt.xlabel('Комбинация параметров', fontsize=14)
+ plt.ylabel('Значение метрики', fontsize=14)
+
+ # Поворачиваем подписи по оси X для лучшей читаемости
+ plt.xticks(rotation=45, ha='right', fontsize=10)
+
+ # Настраиваем легенду
+ plt.legend(title='Метрика', fontsize=12)
+
+ # Добавляем сетку
+ plt.grid(axis='y', linestyle='--', alpha=0.7)
+
+ # Настраиваем макет
+ plt.tight_layout()
+
+ # Сохраняем график
+ file_name = f"best_combinations_comparison.png"
+ plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
+ plt.close()
+
+ print(f"Создан график сравнения лучших комбинаций: {file_name}")
+
+
+def generate_plots(combined_results: pd.DataFrame, macro_metrics: pd.DataFrame, plots_dir: str) -> None:
+ """
+ Генерирует набор графиков с помощью seaborn и сохраняет их в указанную директорию.
+ Фокусируется в первую очередь на recall-метриках как наиболее важных.
+
+ Args:
+ combined_results: DataFrame с объединенными результатами (weighted метрики)
+ macro_metrics: DataFrame с macro метриками
+ plots_dir: Директория для сохранения графиков
+ """
+ # Создаем директорию для графиков, если она не существует
+ setup_plot_directory(plots_dir)
+
+ # Настраиваем стиль для графиков
+ sns.set_style("whitegrid")
+ plt.rcParams['font.family'] = 'DejaVu Sans'
+
+ # Получаем список моделей для построения графиков
+ models = combined_results['model'].unique()
+ top_n_values = [10, 20, 50, 100]
+
+ print(f"Генерация графиков для {len(models)} моделей...")
+
+ # 0. Добавляем анализ наилучших комбинаций параметров
+ # Определяем метрики для анализа - фокусируемся на recall
+ weighted_recall_metrics = ['text_recall', 'doc_recall']
+
+ # Находим лучшие комбинации параметров
+ best_combinations = find_best_combinations(combined_results, weighted_recall_metrics)
+
+ # Создаем график сравнения лучших комбинаций
+ if not best_combinations.empty:
+ create_best_combinations_plot(best_combinations, weighted_recall_metrics, plots_dir)
+
+ # Если доступны macro метрики, делаем то же самое для них
+ if macro_metrics is not None:
+ macro_recall_metrics = ['macro_text_recall', 'macro_doc_recall']
+ macro_best_combinations = find_best_combinations(macro_metrics, macro_recall_metrics)
+
+ if not macro_best_combinations.empty:
+ create_best_combinations_plot(macro_best_combinations, macro_recall_metrics, plots_dir)
+
+ # 1. Создаем графики сравнения моделей для weighted метрик
+ # Фокусируемся на recall-метриках
+ weighted_metrics = {
+ 'text': ['text_recall'], # Только text_recall
+ 'doc': ['doc_recall'] # Только doc_recall
+ }
+
+ for top_n in top_n_values:
+ for metrics_group, metrics in weighted_metrics.items():
+ create_model_comparison_plot(combined_results, metrics, top_n, plots_dir)
+
+ # 2. Если доступны macro метрики, создаем графики на их основе
+ if macro_metrics is not None:
+ print("Создание графиков на основе macro метрик...")
+ macro_metrics_groups = {
+ 'text': ['macro_text_recall'], # Только macro_text_recall
+ 'doc': ['macro_doc_recall'] # Только macro_doc_recall
+ }
+
+ for top_n in top_n_values:
+ for metrics_group, metrics in macro_metrics_groups.items():
+ create_model_comparison_plot(macro_metrics, metrics, top_n, plots_dir)
+
+ # 3. Создаем графики зависимости от top_n
+ for metrics_type, df in [("weighted", combined_results), ("macro", macro_metrics)]:
+ if df is None:
+ continue
+
+ metrics_to_plot = []
+ if metrics_type == "weighted":
+ metrics_to_plot = ['text_recall', 'doc_recall'] # Только recall-метрики
+ else:
+ metrics_to_plot = ['macro_text_recall', 'macro_doc_recall'] # Только macro recall-метрики
+
+ for metric in metrics_to_plot:
+ create_top_n_plot(df, models, metric, plots_dir)
+
+ # 4. Для каждой модели создаем графики по размеру чанка
+ for model in models:
+ # Выбираем 2 значения top_n для анализа
+ for top_n in [20, 50]:
+ # Создаем графики с recall-метриками
+ weighted_metrics_to_combine = ['text_recall']
+ create_chunk_size_plot(combined_results, model, weighted_metrics_to_combine, top_n, plots_dir)
+
+ doc_metrics_to_combine = ['doc_recall']
+ create_chunk_size_plot(combined_results, model, doc_metrics_to_combine, top_n, plots_dir)
+
+ # Если есть macro метрики, создаем соответствующие графики
+ if macro_metrics is not None:
+ macro_metrics_to_combine = ['macro_text_recall']
+ create_chunk_size_plot(macro_metrics, model, macro_metrics_to_combine, top_n, plots_dir)
+
+ macro_doc_metrics_to_combine = ['macro_doc_recall']
+ create_chunk_size_plot(macro_metrics, model, macro_doc_metrics_to_combine, top_n, plots_dir)
+
+ # 5. Создаем тепловые карты для моделей
+ for top_n in [20, 50]:
+ for metric_prefix in ["", "macro_"]:
+ for metric_type in ["text_recall", "doc_recall"]:
+ metric = f"{metric_prefix}{metric_type}"
+ # Используем соответствующий DataFrame
+ if metric_prefix and macro_metrics is None:
+ continue
+ df_to_use = macro_metrics if metric_prefix else combined_results
+ create_heatmap(df_to_use, models, metric, top_n, plots_dir)
+
+ print(f"Создание графиков завершено в директории {plots_dir}")
+
+
+def print_best_combinations(best_df: pd.DataFrame) -> None:
+ """
+ Выводит информацию о лучших комбинациях параметров.
+
+ Args:
+ best_df: DataFrame с лучшими комбинациями
+ """
+ if best_df.empty:
+ print("Не найдено лучших комбинаций")
+ return
+
+ print("\n=== ЛУЧШИЕ КОМБИНАЦИИ ПАРАМЕТРОВ ===")
+
+ # Выводим абсолютно лучшую комбинацию, если она есть
+ absolute_best = best_df[best_df.get('absolute_best', False) == True]
+ if not absolute_best.empty:
+ row = absolute_best.iloc[0]
+ print(f"\nАБСОЛЮТНО ЛУЧШАЯ КОМБИНАЦИЯ:")
+ print(f" Модель: {row['model']}")
+ print(f" Размер чанка: {row['words_per_chunk']} слов")
+ print(f" Перекрытие: {row['overlap_words']} слов ({row['overlap_percentage']}%)")
+ print(f" top_n: {row['top_n']}")
+
+ # Выводим значения метрик
+ recall_metrics = [col for col in best_df.columns if 'recall' in col and col != 'combined_recall']
+ for metric in recall_metrics:
+ print(f" {metric}: {row[metric]:.4f}")
+
+ print("\n=== ТОП-5 ЛУЧШИХ КОМБИНАЦИЙ ===")
+ for i, row in best_df.head(5).iterrows():
+ print(f"\n#{i+1}: {row['model']}, w={row['words_per_chunk']}, o={row['overlap_words']}, top_n={row['top_n']}")
+
+ # Выводим значения метрик
+ recall_metrics = [col for col in best_df.columns if 'recall' in col and col != 'combined_recall']
+ for metric in recall_metrics:
+ print(f" {metric}: {row[metric]:.4f}")
+
+ print("\n=======================================")
+
+
+def create_combined_excel(combined_results: pd.DataFrame, question_metrics: pd.DataFrame,
+ macro_metrics: pd.DataFrame = None, output_file: str = "combined_results.xlsx") -> None:
+ """
+ Создает Excel-файл с несколькими листами, содержащими различные срезы данных.
+ Добавляет автофильтры и применяет форматирование.
+
+ Args:
+ combined_results: DataFrame с объединенными результатами
+ question_metrics: DataFrame с метриками по вопросам
+ macro_metrics: DataFrame с macro метриками (опционально)
+ output_file: Путь к выходному Excel-файлу
+ """
+ print(f"Создание Excel-файла {output_file}...")
+
+ # Создаем новый Excel-файл
+ workbook = Workbook()
+
+ # Удаляем стандартный лист
+ default_sheet = workbook.active
+ workbook.remove(default_sheet)
+
+ # Подготавливаем данные для различных листов
+ sheets_data = {
+ "Исходные данные": combined_results,
+ "Сводка по моделям": prepare_summary_by_model_top_n(combined_results, macro_metrics),
+ "Сводка по чанкингу": prepare_summary_by_chunking_params_top_n(combined_results, macro_metrics),
+ "Лучшие конфигурации": prepare_best_configurations(combined_results, macro_metrics)
+ }
+
+ # Если есть метрики по вопросам, добавляем лист с ними
+ if question_metrics is not None:
+ sheets_data["Метрики по вопросам"] = question_metrics
+
+ # Если есть macro метрики, добавляем лист с ними
+ if macro_metrics is not None:
+ sheets_data["Macro метрики"] = macro_metrics
+
+ # Создаем листы и добавляем данные
+ for sheet_name, data in sheets_data.items():
+ if data is not None and not data.empty:
+ sheet = workbook.create_sheet(title=sheet_name)
+ for r in dataframe_to_rows(data, index=False, header=True):
+ sheet.append(r)
+
+ # Применяем форматирование
+ apply_formatting(workbook)
+
+ # Сохраняем файл
+ workbook.save(output_file)
+ print(f"Excel-файл создан: {output_file}")
+
+
+def calculate_macro_metrics(question_metrics: pd.DataFrame) -> pd.DataFrame:
+ """
+ Вычисляет macro метрики на основе результатов по вопросам.
+
+ Args:
+ question_metrics: DataFrame с метриками по вопросам
+
+ Returns:
+ DataFrame с macro метриками
+ """
+ if question_metrics is None:
+ return None
+
+ print("Вычисление macro метрик на основе метрик по вопросам...")
+
+ # Группируем по конфигурации (модель, параметры чанкинга, top_n)
+ grouped_metrics = question_metrics.groupby(['model', 'words_per_chunk', 'overlap_words', 'top_n'])
+
+ # Для каждой группы вычисляем среднее значение метрик (macro)
+ macro_metrics = grouped_metrics.agg({
+ 'text_precision': 'mean', # Macro precision = среднее precision по всем вопросам
+ 'text_recall': 'mean', # Macro recall = среднее recall по всем вопросам
+ 'text_f1': 'mean', # Macro F1 = среднее F1 по всем вопросам
+ 'doc_precision': 'mean',
+ 'doc_recall': 'mean',
+ 'doc_f1': 'mean'
+ }).reset_index()
+
+ # Добавляем префикс "macro_" к названиям метрик для ясности
+ for col in ['text_precision', 'text_recall', 'text_f1', 'doc_precision', 'doc_recall', 'doc_f1']:
+ macro_metrics.rename(columns={col: f'macro_{col}'}, inplace=True)
+
+ # Добавляем процент перекрытия
+ macro_metrics['overlap_percentage'] = (macro_metrics['overlap_words'] / macro_metrics['words_per_chunk'] * 100).round(1)
+
+ print(f"Вычислено {len(macro_metrics)} наборов macro метрик")
+
+ return macro_metrics
+
+
+def main():
+ """Основная функция скрипта."""
+ args = parse_args()
+
+ # Загружаем результаты из CSV-файлов
+ combined_results = load_results_files(args.results_dir)
+
+ # Загружаем метрики по вопросам (если есть)
+ question_metrics = load_question_metrics_files(args.results_dir)
+
+ # Вычисляем macro метрики на основе метрик по вопросам
+ macro_metrics = calculate_macro_metrics(question_metrics)
+
+ # Находим лучшие комбинации параметров
+ best_combinations_weighted = find_best_combinations(combined_results, ['text_recall', 'doc_recall'])
+ print_best_combinations(best_combinations_weighted)
+
+ if macro_metrics is not None:
+ best_combinations_macro = find_best_combinations(macro_metrics, ['macro_text_recall', 'macro_doc_recall'])
+ print_best_combinations(best_combinations_macro)
+
+ # Создаем объединенный Excel-файл с данными
+ create_combined_excel(combined_results, question_metrics, macro_metrics, args.output_file)
+
+ # Генерируем графики с помощью seaborn
+ print(f"Генерация графиков и сохранение их в директорию: {args.plots_dir}")
+ generate_plots(combined_results, macro_metrics, args.plots_dir)
+
+ print("Готово! Результаты сохранены в Excel и графики созданы.")
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/lib/extractor/scripts/debug_question_chunks.py b/lib/extractor/scripts/debug_question_chunks.py
new file mode 100644
index 0000000000000000000000000000000000000000..98af762bb5707606789da7501cc46b2c07075aaa
--- /dev/null
+++ b/lib/extractor/scripts/debug_question_chunks.py
@@ -0,0 +1,392 @@
+#!/usr/bin/env python
+"""
+Скрипт для отладки и анализа чанков, найденных для конкретного вопроса.
+Показывает, какие чанки находятся, какие пункты ожидаются и значения метрик нечеткого сравнения.
+"""
+
+import argparse
+import json
+import os
+import sys
+from difflib import SequenceMatcher
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from sklearn.metrics.pairwise import cosine_similarity
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+
+# Константы для настройки
+DATA_FOLDER = "data/docs" # Путь к папке с документами
+MODEL_NAME = "intfloat/e5-base" # Название модели для векторизации
+DATASET_PATH = "data/dataset.xlsx" # Путь к Excel-датасету с вопросами
+OUTPUT_DIR = "data" # Директория для сохранения результатов
+TOP_N_VALUES = [5, 10, 20, 30, 50, 100] # Значения N для анализа
+THRESHOLD = 0.6
+
+
+def parse_args():
+ """
+ Парсит аргументы командной строки.
+
+ Returns:
+ Аргументы командной строки
+ """
+ parser = argparse.ArgumentParser(description="Скрипт для отладки чанкинга на конкретном вопросе")
+
+ parser.add_argument("--data-folder", type=str, default=DATA_FOLDER,
+ help=f"Путь к папке с документами (по умолчанию: {DATA_FOLDER})")
+ parser.add_argument("--model-name", type=str, default=MODEL_NAME,
+ help=f"Название модели для векторизации (по умолчанию: {MODEL_NAME})")
+ parser.add_argument("--dataset-path", type=str, default=DATASET_PATH,
+ help=f"Путь к Excel-датасету с вопросами (по умолчанию: {DATASET_PATH})")
+ parser.add_argument("--output-dir", type=str, default=OUTPUT_DIR,
+ help=f"Директория для сохранения результатов (по умолчанию: {OUTPUT_DIR})")
+ parser.add_argument("--question-id", type=int, required=True,
+ help="ID вопроса для отладки")
+ parser.add_argument("--top-n", type=int, default=20,
+ help="Количество чанков в топе для отладки (по умолчанию: 20)")
+ parser.add_argument("--words-per-chunk", type=int, default=50,
+ help="Количество слов в чанке для fixed_size стратегии (по умолчанию: 50)")
+ parser.add_argument("--overlap-words", type=int, default=25,
+ help="Количество слов перекрытия для fixed_size стратегии (по умолчанию: 25)")
+
+ return parser.parse_args()
+
+
+def load_questions_dataset(file_path: str) -> pd.DataFrame:
+ """
+ Загружает датасет с вопросами из Excel-файла.
+
+ Args:
+ file_path: Путь к Excel-файлу
+
+ Returns:
+ DataFrame с вопросами и пунктами
+ """
+ print(f"Загрузка датасета из {file_path}...")
+
+ df = pd.read_excel(file_path)
+ print(f"Загружен датасет со столбцами: {df.columns.tolist()}")
+
+ # Преобразуем NaN в пустые строки для текстовых полей
+ text_columns = ['question', 'text', 'item_type']
+ for col in text_columns:
+ if col in df.columns:
+ df[col] = df[col].fillna('')
+
+ return df
+
+
+def load_embeddings_and_data(filename: str, output_dir: str) -> tuple[np.ndarray | None, pd.DataFrame | None]:
+ """
+ Загружает эмбеддинги и соответствующие данные из файлов.
+
+ Args:
+ filename: Базовое имя файла
+ output_dir: Директория, где хранятся файлы
+
+ Returns:
+ Кортеж (эмбеддинги, данные) или (None, None), если файлы не найдены
+ """
+ embeddings_path = os.path.join(output_dir, f"{filename}_embeddings.npy")
+ data_path = os.path.join(output_dir, f"{filename}_data.csv")
+
+ if os.path.exists(embeddings_path) and os.path.exists(data_path):
+ print(f"Загрузка данных из {embeddings_path} и {data_path}...")
+ embeddings = np.load(embeddings_path)
+ data = pd.read_csv(data_path)
+ return embeddings, data
+
+ print(f"Ошибка: файлы {embeddings_path} и {data_path} не найдены.")
+ print("Сначала запустите скрипт evaluate_chunking.py для создания эмбеддингов.")
+ sys.exit(1)
+
+
+def calculate_chunk_overlap(chunk_text: str, punct_text: str) -> float:
+ """
+ Рассчитывает степень перекрытия между чанком и пунктом.
+
+ Args:
+ chunk_text: Текст чанка
+ punct_text: Текст пункта
+
+ Returns:
+ Коэффициент перекрытия от 0 до 1
+ """
+ # Если чанк входит в пункт, возвращаем 1.0 (полное вхождение)
+ if chunk_text in punct_text:
+ return 1.0
+
+ # Если пункт входит в чанк, возвращаем соотношение длин
+ if punct_text in chunk_text:
+ return len(punct_text) / len(chunk_text)
+
+ # Используем SequenceMatcher для нечеткого сравнения
+ matcher = SequenceMatcher(None, chunk_text, punct_text)
+
+ # Находим наибольшую общую подстроку
+ match = matcher.find_longest_match(0, len(chunk_text), 0, len(punct_text))
+
+ # Если совпадений нет
+ if match.size == 0:
+ return 0.0
+
+ # Возвращаем соотношение длины совпадения к минимальной длине
+ return match.size / min(len(chunk_text), len(punct_text))
+
+
+def format_text_for_display(text: str, max_length: int = 100) -> str:
+ """
+ Форматирует текст для отображения, обрезая его при необходимости.
+
+ Args:
+ text: Исходный текст
+ max_length: Максимальная длина для отображения
+
+ Returns:
+ Отформатированный текст
+ """
+ if len(text) <= max_length:
+ return text
+ return text[:max_length] + "..."
+
+
+def analyze_question(
+ question_id: int,
+ questions_df: pd.DataFrame,
+ chunks_df: pd.DataFrame,
+ question_embeddings: np.ndarray,
+ chunk_embeddings: np.ndarray,
+ question_id_to_idx: dict,
+ top_n: int
+) -> dict:
+ """
+ Анализирует конкретный вопрос и его релевантные чанки.
+
+ Args:
+ question_id: ID вопроса для анализа
+ questions_df: DataFrame с вопросами
+ chunks_df: DataFrame с чанками
+ question_embeddings: Эмбеддинги вопросов
+ chunk_embeddings: Эмбеддинги чанков
+ question_id_to_idx: Словарь соответствия ID вопроса и его индекса
+ top_n: Количество чанков в топе
+
+ Returns:
+ Словарь с результатами анализа
+ """
+ # Проверяем, есть ли вопрос с таким ID
+ if question_id not in question_id_to_idx:
+ print(f"Ошибка: вопрос с ID {question_id} не найден в данных")
+ sys.exit(1)
+
+ # Получаем строки для выбранного вопроса
+ question_rows = questions_df[questions_df['id'] == question_id]
+ if len(question_rows) == 0:
+ print(f"Ошибка: вопрос с ID {question_id} не найден в исходном датасете")
+ sys.exit(1)
+
+ # Получаем текст вопроса и его индекс в массиве эмбеддингов
+ question_text = question_rows['question'].iloc[0]
+ question_idx = question_id_to_idx[question_id]
+
+ # Получаем ожидаемые пункты для вопроса
+ expected_puncts = question_rows['text'].tolist()
+
+ # Вычисляем косинусную близость между вопросом и всеми чанками
+ similarity = cosine_similarity([question_embeddings[question_idx]], chunk_embeddings)[0]
+
+ # Получаем связанные документы, если есть
+ related_docs = []
+ if 'filename' in question_rows.columns:
+ related_docs = question_rows['filename'].unique().tolist()
+ related_docs = [doc for doc in related_docs if doc and not pd.isna(doc)]
+
+ # Результаты для всех документов
+ all_results = []
+
+ # Обрабатываем каждый связанный документ
+ if related_docs:
+ for doc_name in related_docs:
+ # Фильтруем чанки по имени документа
+ doc_chunks = chunks_df[chunks_df['doc_name'] == doc_name]
+ if doc_chunks.empty:
+ continue
+
+ # Индексы чанков для документа
+ doc_chunk_indices = doc_chunks.index.tolist()
+
+ # Получаем значения близости для чанков документа
+ doc_similarities = [similarity[chunks_df.index.get_loc(idx)] for idx in doc_chunk_indices]
+
+ # Создаем словарь индекс -> схожесть
+ similarity_dict = {idx: sim for idx, sim in zip(doc_chunk_indices, doc_similarities)}
+
+ # Сортируем индексы по убыванию похожести
+ sorted_indices = sorted(similarity_dict.keys(), key=lambda x: similarity_dict[x], reverse=True)
+
+ # Берем топ-N
+ top_indices = sorted_indices[:min(top_n, len(sorted_indices))]
+
+ # Получаем топ-N чанков
+ top_chunks = chunks_df.iloc[top_indices]
+
+ # Формируем результаты для документа
+ doc_results = {
+ 'doc_name': doc_name,
+ 'top_chunks': []
+ }
+
+ # Для каждого чанка
+ for idx, chunk in top_chunks.iterrows():
+ # Вычисляем перекрытие с каждым пунктом
+ overlaps = []
+ for punct in expected_puncts:
+ overlap = calculate_chunk_overlap(chunk['text'], punct)
+ overlaps.append({
+ 'punct': format_text_for_display(punct),
+ 'overlap': overlap
+ })
+
+ # Находим максимальное перекрытие
+ max_overlap = max(overlaps, key=lambda x: x['overlap']) if overlaps else {'overlap': 0}
+
+ # Добавляем в результаты
+ doc_results['top_chunks'].append({
+ 'chunk_id': chunk['id'],
+ 'chunk_text': format_text_for_display(chunk['text']),
+ 'similarity': similarity_dict[idx],
+ 'overlaps': overlaps,
+ 'max_overlap': max_overlap['overlap'],
+ 'is_relevant': max_overlap['overlap'] >= THRESHOLD # Используем порог 0.7
+ })
+
+ all_results.append(doc_results)
+ else:
+ # Если нет связанных документов, анализируем чанки из всех документов
+ # Получаем индексы для топ-N чанков по близости
+ top_indices = np.argsort(similarity)[-top_n:][::-1]
+
+ # Получаем топ-N чанков
+ top_chunks = chunks_df.iloc[top_indices]
+
+ # Группируем чанки по документам
+ doc_groups = top_chunks.groupby('doc_name')
+
+ for doc_name, group in doc_groups:
+ doc_results = {
+ 'doc_name': doc_name,
+ 'top_chunks': []
+ }
+
+ for idx, chunk in group.iterrows():
+ # Вычисляем перекрытие с каждым пунктом
+ overlaps = []
+ for punct in expected_puncts:
+ overlap = calculate_chunk_overlap(chunk['text'], punct)
+ overlaps.append({
+ 'punct': format_text_for_display(punct),
+ 'overlap': overlap
+ })
+
+ # Находим максимальное перекрытие
+ max_overlap = max(overlaps, key=lambda x: x['overlap']) if overlaps else {'overlap': 0}
+
+ # Добавляем в результаты
+ doc_results['top_chunks'].append({
+ 'chunk_id': chunk['id'],
+ 'chunk_text': format_text_for_display(chunk['text']),
+ 'similarity': similarity[chunks_df.index.get_loc(idx)],
+ 'overlaps': overlaps,
+ 'max_overlap': max_overlap['overlap'],
+ 'is_relevant': max_overlap['overlap'] >= THRESHOLD # Используем порог 0.7
+ })
+
+ all_results.append(doc_results)
+
+ # Формируем общие результаты для вопроса
+ results = {
+ 'question_id': question_id,
+ 'question_text': question_text,
+ 'expected_puncts': [format_text_for_display(punct) for punct in expected_puncts],
+ 'related_docs': related_docs,
+ 'results_by_doc': all_results
+ }
+
+ return results
+
+
+def main():
+ """
+ Основная функция скрипта.
+ """
+ args = parse_args()
+
+ # Загружаем датасет с вопросами
+ questions_df = load_questions_dataset(args.dataset_path)
+
+ # Формируем уникальное имя для сохраненных файлов на основе параметров стратегии и модели
+ strategy_config_str = f"fixed_size_w{args.words_per_chunk}_o{args.overlap_words}"
+ chunks_filename = f"chunks_{strategy_config_str}_{args.model_name.replace('/', '_')}"
+ questions_filename = f"questions_{args.model_name.replace('/', '_')}"
+
+ # Загружаем сохраненные эмбеддинги и данные
+ chunk_embeddings, chunks_df = load_embeddings_and_data(chunks_filename, args.output_dir)
+ question_embeddings, questions_df_with_embeddings = load_embeddings_and_data(questions_filename, args.output_dir)
+
+ # Создаем словарь соответствия id вопроса и его индекса в эмбеддингах
+ question_id_to_idx = {
+ int(row['id']): i
+ for i, (_, row) in enumerate(questions_df_with_embeddings.iterrows())
+ }
+
+ # Анализируем выбранный вопрос для указанного top_n
+ results = analyze_question(
+ args.question_id,
+ questions_df,
+ chunks_df,
+ question_embeddings,
+ chunk_embeddings,
+ question_id_to_idx,
+ args.top_n
+ )
+
+ # Сохраняем результаты в JSON файл
+ output_filename = f"debug_question_{args.question_id}_top{args.top_n}.json"
+ output_path = os.path.join(args.output_dir, output_filename)
+
+ with open(output_path, 'w', encoding='utf-8') as f:
+ json.dump(results, f, ensure_ascii=False, indent=2)
+
+ print(f"Результаты сохранены в {output_path}")
+
+ # Выводим краткую информацию
+ print(f"\nАнализ вопроса ID {args.question_id}: {results['question_text']}")
+ print(f"Ожидаемые пункты: {len(results['expected_puncts'])}")
+ print(f"Связанные документы: {results['related_docs']}")
+
+ # Статистика релевантности
+ relevant_chunks = 0
+ total_chunks = 0
+
+ for doc_result in results['results_by_doc']:
+ doc_relevant = sum(1 for chunk in doc_result['top_chunks'] if chunk['is_relevant'])
+ doc_total = len(doc_result['top_chunks'])
+
+ print(f"\nДокумент: {doc_result['doc_name']}")
+ print(f"Релевантных чанков: {doc_relevant} из {doc_total} ({doc_relevant/doc_total*100:.1f}%)")
+
+ relevant_chunks += doc_relevant
+ total_chunks += doc_total
+
+ if total_chunks > 0:
+ print(f"\nОбщая точность: {relevant_chunks/total_chunks*100:.1f}%")
+ else:
+ print("\nНе найдено чанков для анализа")
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/lib/extractor/scripts/evaluate_chunking.py b/lib/extractor/scripts/evaluate_chunking.py
new file mode 100644
index 0000000000000000000000000000000000000000..3f2061cbf76651379fa68d8adb893bc127f02a50
--- /dev/null
+++ b/lib/extractor/scripts/evaluate_chunking.py
@@ -0,0 +1,800 @@
+#!/usr/bin/env python
+"""
+Скрипт для оценки качества различных стратегий чанкинга.
+Сравнивает стратегии на основе релевантности чанков к вопросам.
+"""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from fuzzywuzzy import fuzz
+from sklearn.metrics.pairwise import cosine_similarity
+from tqdm import tqdm
+from transformers import AutoModel, AutoTokenizer
+
+# Константы для настройки
+DATA_FOLDER = "data/docs" # Путь к папке с документами
+MODEL_NAME = "intfloat/e5-base" # Название модели для векторизации
+DATASET_PATH = "data/dataset.xlsx" # Путь к Excel-датасету с вопросами
+BATCH_SIZE = 8 # Размер батча для векторизации
+DEVICE = "cuda:1" if torch.cuda.is_available() else "cpu" # Устройство для вычислений
+SIMILARITY_THRESHOLD = 0.7 # Порог для нечеткого сравнения
+OUTPUT_DIR = "data" # Директория для сохранения результатов
+TOP_CHUNKS_DIR = "data/top_chunks" # Директория для сохранения топ-чанков
+TOP_N_VALUES = [5, 10, 20, 30, 50, 70, 100] # Значения N для оценки
+
+# Параметры стратегий чанкинга
+FIXED_SIZE_CONFIG = {
+ "words_per_chunk": 50, # Количество слов в чанке
+ "overlap_words": 25 # Количество слов перекрытия
+}
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from ntr_fileparser import UniversalParser
+
+from ntr_text_fragmentation import Destructurer
+
+
+def _average_pool(
+ last_hidden_states: torch.Tensor, attention_mask: torch.Tensor
+ ) -> torch.Tensor:
+ """
+ Расчёт усредненного эмбеддинга по всем токенам
+
+ Args:
+ last_hidden_states: Матрица эмбеддингов отдельных токенов размерности (batch_size, seq_len, embedding_size) - последний скрытый слой
+ attention_mask: Маска, чтобы не учитывать при усреднении пустые токены
+
+ Returns:
+ torch.Tensor - Усредненный эмбеддинг размерности (batch_size, embedding_size)
+ """
+ last_hidden = last_hidden_states.masked_fill(
+ ~attention_mask[..., None].bool(), 0.0
+ )
+ return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+
+def parse_args():
+ """
+ Парсит аргументы командной строки.
+
+ Returns:
+ Аргументы командной строки
+ """
+ parser = argparse.ArgumentParser(description="Скрипт для оценки качества чанкинга")
+
+ parser.add_argument("--data-folder", type=str, default=DATA_FOLDER,
+ help=f"Путь к папке с документами (по умолчанию: {DATA_FOLDER})")
+ parser.add_argument("--model-name", type=str, default=MODEL_NAME,
+ help=f"Название модели для векторизации (по умолчанию: {MODEL_NAME})")
+ parser.add_argument("--dataset-path", type=str, default=DATASET_PATH,
+ help=f"Путь к Excel-датасету с вопросами (по умолчанию: {DATASET_PATH})")
+ parser.add_argument("--batch-size", type=int, default=BATCH_SIZE,
+ help=f"Размер батча для векторизации (по умолчанию: {BATCH_SIZE})")
+ parser.add_argument("--similarity-threshold", type=float, default=SIMILARITY_THRESHOLD,
+ help=f"Порог для нечеткого сравнения (по умолчанию: {SIMILARITY_THRESHOLD})")
+ parser.add_argument("--output-dir", type=str, default=OUTPUT_DIR,
+ help=f"Директория для сохранения результатов (по умолчанию: {OUTPUT_DIR})")
+ parser.add_argument("--force-recompute", action="store_true",
+ help="Принудительно пересчитать эмбеддинги, игнорируя сохраненные")
+ parser.add_argument("--use-sentence-transformers", action="store_true",
+ help="Использовать библиотеку sentence_transformers для извлечения эмбеддингов (для FRIDA и других моделей)")
+ parser.add_argument("--device", type=str, default=DEVICE,
+ help=f"Устройство для вычислений (по умолчанию: {DEVICE})")
+
+ # Параметры для fixed_size стратегии
+ parser.add_argument("--words-per-chunk", type=int, default=FIXED_SIZE_CONFIG["words_per_chunk"],
+ help=f"Количество слов в чанке для fixed_size стратегии (по умолчанию: {FIXED_SIZE_CONFIG['words_per_chunk']})")
+ parser.add_argument("--overlap-words", type=int, default=FIXED_SIZE_CONFIG["overlap_words"],
+ help=f"Количество слов перекрытия для fixed_size стратегии (по умолчанию: {FIXED_SIZE_CONFIG['overlap_words']})")
+
+ return parser.parse_args()
+
+
+def read_documents(folder_path: str) -> dict:
+ """
+ Читает все документы из указанной папки.
+
+ Args:
+ folder_path: Путь к папке с документами
+
+ Returns:
+ Словарь {имя_файла: parsed_document}
+ """
+ print(f"Чтение документов из {folder_path}...")
+ parser = UniversalParser()
+ documents = {}
+
+ for file_path in tqdm(list(Path(folder_path).glob("*.docx")), desc="Чтение документов"):
+ try:
+ doc_name = file_path.stem
+ documents[doc_name] = parser.parse_by_path(str(file_path))
+ except Exception as e:
+ print(f"Ошибка при чтении файла {file_path}: {e}")
+
+ return documents
+
+
+def process_documents(documents: dict, fixed_size_config: dict) -> pd.DataFrame:
+ """
+ Обрабатывает документы со стратегией fixed_size для чанкинга.
+
+ Args:
+ documents: Словарь с распарсенными документами
+ fixed_size_config: Конфигурация для fixed_size стратегии
+
+ Returns:
+ DataFrame с чанками
+ """
+ print("Обработка документов стратегией fixed_size...")
+
+ all_data = []
+
+ for doc_name, document in tqdm(documents.items(), desc="Применение стратегии fixed_size"):
+ # Стратегия fixed_size для чанкинга
+ destructurer = Destructurer(document)
+ destructurer.configure('fixed_size',
+ words_per_chunk=fixed_size_config["words_per_chunk"],
+ overlap_words=fixed_size_config["overlap_words"])
+ fixed_size_entities, _ = destructurer.destructure()
+
+ # Обрабатываем только сущности для поиска
+ for entity in fixed_size_entities:
+ if hasattr(entity, 'use_in_search') and entity.use_in_search:
+ entity_data = {
+ 'id': str(entity.id),
+ 'doc_name': doc_name,
+ 'name': entity.name,
+ 'text': entity.text,
+ 'type': entity.type,
+ 'strategy': 'fixed_size',
+ 'metadata': json.dumps(entity.metadata, ensure_ascii=False)
+ }
+ all_data.append(entity_data)
+
+ # Создаем DataFrame
+ df = pd.DataFrame(all_data)
+
+ # Фильтруем по типу, исключая Document
+ df = df[df['type'] != 'Document']
+
+ return df
+
+
+def load_questions_dataset(file_path: str) -> pd.DataFrame:
+ """
+ Загружает датасет с вопросами из Excel-файла.
+
+ Args:
+ file_path: Путь к Excel-файлу
+
+ Returns:
+ DataFrame с вопросами и пунктами
+ """
+ print(f"Загрузка датасета из {file_path}...")
+
+ df = pd.read_excel(file_path)
+ print(f"Загружен датасет со столбцами: {df.columns.tolist()}")
+
+ # Преобразуем NaN в пустые строки для текстовых полей
+ text_columns = ['question', 'text', 'item_type']
+ for col in text_columns:
+ if col in df.columns:
+ df[col] = df[col].fillna('')
+
+ return df
+
+
+def setup_model_and_tokenizer(model_name: str, use_sentence_transformers: bool = False, device: str = DEVICE):
+ """
+ Инициализирует модель и токенизатор.
+
+ Args:
+ model_name: Название предобученной модели
+ use_sentence_transformers: Использовать ли библиотеку sentence_transformers
+ device: Устройство для вычислений
+
+ Returns:
+ Кортеж (модель, токенизатор) или объект SentenceTransformer
+ """
+ print(f"Загрузка модели {model_name} на устройство {device}...")
+
+ if use_sentence_transformers:
+ try:
+ from sentence_transformers import SentenceTransformer
+ model = SentenceTransformer(model_name, device=device)
+ return model, None
+ except ImportError:
+ print("Библиотека sentence_transformers не установлена. Установите её с помощью pip install sentence-transformers")
+ raise
+ else:
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
+ model = AutoModel.from_pretrained(model_name).to(device)
+ model.eval()
+
+ return model, tokenizer
+
+
+def get_embeddings(texts: list[str], model, tokenizer=None, batch_size: int = BATCH_SIZE, use_sentence_transformers: bool = False, device: str = DEVICE) -> np.ndarray:
+ """
+ Получает эмбеддинги для списка текстов с использованием average pooling или sentence_transformers.
+
+ Args:
+ texts: Список текстов
+ model: Модель для векторизации или SentenceTransformer
+ tokenizer: Токенизатор (None для sentence_transformers)
+ batch_size: Размер батча
+ use_sentence_transformers: Использовать ли библиотеку sentence_transformers
+ device: Устройство для вычислений
+
+ Returns:
+ Массив эмбеддингов
+ """
+ if use_sentence_transformers:
+ # Используем sentence_transformers для получения эмбеддингов
+ all_embeddings = []
+
+ for i in tqdm(range(0, len(texts), batch_size), desc="Векторизация текстов (sentence_transformers)"):
+ batch_texts = texts[i:i+batch_size]
+
+ # Получаем эмбеддинги с помощью sentence_transformers
+ embeddings = model.encode(batch_texts, batch_size=batch_size, show_progress_bar=False)
+ all_embeddings.append(embeddings)
+
+ return np.vstack(all_embeddings)
+ else:
+ # Используем стандартный подход с average pooling
+ all_embeddings = []
+
+ for i in tqdm(range(0, len(texts), batch_size), desc="Векторизация текстов"):
+ batch_texts = texts[i:i+batch_size]
+
+ # Токенизация с обрезкой и padding
+ encoding = tokenizer(
+ batch_texts,
+ padding=True,
+ truncation=True,
+ max_length=512,
+ return_tensors="pt"
+ ).to(device)
+
+ # Получаем эмбеддинги с average pooling
+ with torch.no_grad():
+ outputs = model(**encoding)
+ embeddings = _average_pool(outputs.last_hidden_state, encoding["attention_mask"])
+ all_embeddings.append(embeddings.cpu().numpy())
+
+ return np.vstack(all_embeddings)
+
+
+def calculate_chunk_overlap(chunk_text: str, punct_text: str) -> float:
+ """
+ Рассчитывает степень перекрытия между чанком и пунктом с использованием partial_ratio.
+
+ Args:
+ chunk_text: Текст чанка
+ punct_text: Текст пункта
+
+ Returns:
+ Коэффициент перекрытия от 0 до 1
+ """
+ # Если чанк входит в пункт, возвращаем 1.0 (полное вхождение)
+ if chunk_text in punct_text:
+ return 1.0
+
+ # Если пункт входит в чанк, возвращаем соотношение длин
+ if punct_text in chunk_text:
+ return len(punct_text) / len(chunk_text)
+
+ # Используем partial_ratio из fuzzywuzzy, который лучше обрабатывает
+ # случаи, когда один текст является подстрокой другого, даже с небольшими различиями
+ partial_ratio_score = fuzz.partial_ratio(chunk_text, punct_text) / 100.0
+
+ return partial_ratio_score
+
+
+def save_embeddings_and_data(embeddings: np.ndarray, data: pd.DataFrame, filename: str, output_dir: str):
+ """
+ Сохраняет эмбеддинги и соответствующие данные в файлы.
+
+ Args:
+ embeddings: Массив эмбеддингов
+ data: DataFrame с данными
+ filename: Базовое имя файла
+ output_dir: Директория для сохранения
+ """
+ embeddings_path = os.path.join(output_dir, f"{filename}_embeddings.npy")
+ data_path = os.path.join(output_dir, f"{filename}_data.csv")
+
+ # Сохраняем эмбеддинги
+ np.save(embeddings_path, embeddings)
+ print(f"Эмбеддинги сохранены в {embeddings_path}")
+
+ # Сохраняем данные
+ data.to_csv(data_path, index=False)
+ print(f"Данные сохранены в {data_path}")
+
+
+def load_embeddings_and_data(filename: str, output_dir: str) -> tuple[np.ndarray | None, pd.DataFrame | None]:
+ """
+ Загружает эмбеддинги и соответствующие данные из файлов.
+
+ Args:
+ filename: Базовое имя файла
+ output_dir: Директория, где хранятся файлы
+
+ Returns:
+ Кортеж (эмбеддинги, данные) или (None, None), если файлы не найдены
+ """
+ embeddings_path = os.path.join(output_dir, f"{filename}_embeddings.npy")
+ data_path = os.path.join(output_dir, f"{filename}_data.csv")
+
+ if os.path.exists(embeddings_path) and os.path.exists(data_path):
+ print(f"Загрузка данных из {embeddings_path} и {data_path}...")
+ embeddings = np.load(embeddings_path)
+ data = pd.read_csv(data_path)
+ return embeddings, data
+
+ return None, None
+
+
+def save_top_chunks_for_question(
+ question_id: int,
+ question_text: str,
+ question_puncts: list[str],
+ top_chunks: pd.DataFrame,
+ similarities: dict,
+ overlap_data: list,
+ output_dir: str
+):
+ """
+ Сохраняет топ-чанки для конкретного вопроса в JSON-файл.
+
+ Args:
+ question_id: ID вопроса
+ question_text: Текст вопроса
+ question_puncts: Список пунктов, относящихся к вопросу
+ top_chunks: DataFrame с топ-чанками
+ similarities: Словарь с косинусными схожестями для чанков
+ overlap_data: Данные о перекрытии чанков с пунктами
+ output_dir: Директория для сохранения
+ """
+ # Подготавливаем результаты для сохранения
+ chunks_data = []
+
+ for i, (idx, chunk) in enumerate(top_chunks.iterrows()):
+ # Получаем данные о перекрытии для текущего чанка
+ chunk_overlaps = overlap_data[i] if i < len(overlap_data) else []
+
+ # Преобразуем numpy типы в стандартные типы Python
+ similarity = float(similarities.get(idx, 0.0))
+
+ # Формируем данные чанка
+ chunk_data = {
+ 'chunk_id': chunk['id'],
+ 'doc_name': chunk['doc_name'],
+ 'text': chunk['text'],
+ 'similarity': similarity,
+ 'overlaps': chunk_overlaps
+ }
+ chunks_data.append(chunk_data)
+
+ # Преобразуем numpy.int64 в int для question_id
+ question_id = int(question_id)
+
+ # Формируем общий результат
+ result = {
+ 'question_id': question_id,
+ 'question_text': question_text,
+ 'puncts': question_puncts,
+ 'chunks': chunks_data
+ }
+
+ # Создаем имя файла
+ filename = f"question_{question_id}_top_chunks.json"
+ filepath = os.path.join(output_dir, filename)
+
+ # Класс для сериализации numpy типов
+ class NumpyEncoder(json.JSONEncoder):
+ def default(self, obj):
+ if isinstance(obj, np.integer):
+ return int(obj)
+ if isinstance(obj, np.floating):
+ return float(obj)
+ if isinstance(obj, np.ndarray):
+ return obj.tolist()
+ return super().default(obj)
+
+ # Сохраняем в JSON с кастомным энкодером
+ with open(filepath, 'w', encoding='utf-8') as f:
+ json.dump(result, f, ensure_ascii=False, indent=2, cls=NumpyEncoder)
+
+ print(f"Топ-чанки для вопроса {question_id} сохранены в {filepath}")
+
+
+def evaluate_for_top_n_with_mapping(
+ questions_df: pd.DataFrame,
+ chunks_df: pd.DataFrame,
+ question_embeddings: np.ndarray,
+ chunk_embeddings: np.ndarray,
+ question_id_to_idx: dict,
+ top_n: int,
+ similarity_threshold: float,
+ top_chunks_dir: str = None
+) -> tuple[dict[str, float], pd.DataFrame]:
+ """
+ Оценивает качество чанкинга для заданного значения top_n с использованием маппинга id -> индекс.
+
+ Args:
+ questions_df: DataFrame с вопросами и релевантными пунктами (исходный датасет)
+ chunks_df: DataFrame с чанками
+ question_embeddings: Эмбеддинги вопросов
+ chunk_embeddings: Эмбеддинги чанков
+ question_id_to_idx: Словарь соответствия id вопроса и его индекса в массиве эмбеддингов
+ top_n: Количество чанков в топе для каждого вопроса
+ similarity_threshold: Порог для нечеткого сравнения
+ top_chunks_dir: Директория для сохранения топ-чанков (если None, то не сохраняем)
+
+ Returns:
+ Кортеж (словарь с усредненными метриками, DataFrame с метриками по отдельным вопросам)
+ """
+ print(f"Оценка для top-{top_n}...")
+
+ # Вычисляем косинусную близость между вопросами и чанками
+ similarity_matrix = cosine_similarity(question_embeddings, chunk_embeddings)
+
+ # Счетчики для метрик на основе текста
+ total_puncts = 0
+ found_puncts = 0
+ total_chunks = 0
+ relevant_chunks = 0
+
+ # Счетчики для метрик на основе документов
+ total_docs_required = 0
+ found_relevant_docs = 0
+ total_docs_found = 0
+
+ # Для сохранения метрик по отдельным вопросам
+ question_metrics = []
+
+ # Выводим информацию о столбцах для отладки
+ print(f"Столбцы в исходном датасете: {questions_df.columns.tolist()}")
+
+ # Группируем вопросы по id (у нас 20 уникальных вопросов)
+ for question_id in tqdm(questions_df['id'].unique(), desc=f"Оценка top-{top_n}"):
+ # Получаем строки для текущего вопроса из исходного датасета
+ question_rows = questions_df[questions_df['id'] == question_id]
+
+ # Проверяем, есть ли вопрос с таким id в нашем маппинге
+ if question_id not in question_id_to_idx:
+ print(f"Предупреждение: вопрос с id {question_id} отсутствует в маппинге")
+ continue
+
+ # Если нет строк с таким id, пропускаем
+ if len(question_rows) == 0:
+ continue
+
+ # Получаем индекс вопроса в массиве эмбеддингов
+ question_idx = question_id_to_idx[question_id]
+
+ # Получаем текст вопроса
+ question_text = question_rows['question'].iloc[0]
+
+ # Получаем все пункты для этого вопроса
+ puncts = question_rows['text'].tolist()
+ question_total_puncts = len(puncts)
+ total_puncts += question_total_puncts
+
+ # Получаем связанные документы
+ relevant_docs = []
+ if 'filename' in question_rows.columns:
+ relevant_docs = [f for f in question_rows['filename'].unique() if f and not pd.isna(f)]
+ question_total_docs_required = len(relevant_docs)
+ total_docs_required += question_total_docs_required
+ print(f"Найдено {question_total_docs_required} документов для вопроса {question_id}")
+ else:
+ print(f"Столбец 'filename' отсутствует. Используем все документы.")
+ relevant_docs = chunks_df['doc_name'].unique().tolist()
+ question_total_docs_required = len(relevant_docs)
+ total_docs_required += question_total_docs_required
+
+ # Если для вопроса нет релевантных документов, пропускаем
+ if not relevant_docs:
+ print(f"Для вопроса {question_id} нет связанных документов")
+ continue
+
+ # Флаги для отслеживания найденных пунктов
+ punct_found = [False] * question_total_puncts
+
+ # Для отслеживания найденных документов
+ docs_found_for_question = set()
+
+ # Для хранения всех чанков вопроса для ограничения top_n
+ all_question_chunks = []
+ all_question_similarities = []
+
+ # Собираем чанки для всех документов по этому вопросу
+ for filename in relevant_docs:
+ if not filename or pd.isna(filename):
+ continue
+
+ # Фильтруем чанки по имени файла
+ doc_chunks = chunks_df[chunks_df['doc_name'] == filename]
+
+ if doc_chunks.empty:
+ print(f"Предупреждение: документ {filename} не содержит чанков")
+ continue
+
+ # Индексы чанков для текущего файла
+ doc_chunk_indices = doc_chunks.index.tolist()
+
+ # Получаем значения близости для чанков текущего файла
+ doc_similarities = [
+ similarity_matrix[question_idx, chunks_df.index.get_loc(idx)]
+ for idx in doc_chunk_indices
+ ]
+
+ # Добавляем чанки и их схожести к общему списку для вопроса
+ for i, idx in enumerate(doc_chunk_indices):
+ all_question_chunks.append((idx, doc_chunks.iloc[doc_chunks.index.get_indexer([idx])[0]]))
+ all_question_similarities.append(doc_similarities[i])
+
+ # Сортируем все чанки по убыванию схожести и берем top_n
+ sorted_indices = np.argsort(all_question_similarities)[-min(top_n, len(all_question_similarities)):][::-1]
+ top_chunks_indices = [all_question_chunks[i][0] for i in sorted_indices]
+ top_chunks = [all_question_chunks[i][1] for i in sorted_indices]
+
+ # Увеличиваем счетчик общего числа чанков
+ question_total_chunks = len(top_chunks)
+ total_chunks += question_total_chunks
+
+ # Для сохранения данных топ-чанков
+ all_top_chunks = pd.DataFrame([chunk for chunk in top_chunks])
+ all_chunk_similarities = {idx: all_question_similarities[i] for i, idx in enumerate([all_question_chunks[j][0] for j in sorted_indices])}
+ all_chunk_overlaps = []
+
+ # Для каждого чанка проверяем его релевантность к пунктам
+ question_relevant_chunks = 0
+
+ for i, chunk in enumerate(top_chunks):
+ is_relevant = False
+ chunk_overlaps = []
+
+ # Добавляем документ в найденные
+ docs_found_for_question.add(chunk['doc_name'])
+
+ # Проверяем перекрытие с каждым пунктом
+ for j, punct in enumerate(puncts):
+ overlap = calculate_chunk_overlap(chunk['text'], punct)
+
+ # Если нужно сохранить топ-чанки и top_n == 20
+ if top_chunks_dir and top_n == 20:
+ chunk_overlaps.append({
+ 'punct_index': j,
+ 'punct_text': punct[:100] + '...' if len(punct) > 100 else punct,
+ 'overlap': overlap
+ })
+
+ # Если перекрытие больше порога, чанк релевантен
+ if overlap >= similarity_threshold:
+ is_relevant = True
+ punct_found[j] = True
+
+ if is_relevant:
+ question_relevant_chunks += 1
+
+ # Если нужно сохранить топ-чанки и top_n == 20
+ if top_chunks_dir and top_n == 20:
+ all_chunk_overlaps.append(chunk_overlaps)
+
+ # Если нужно сохранить топ-чанки и top_n == 20
+ if top_chunks_dir and top_n == 20 and not all_top_chunks.empty:
+ save_top_chunks_for_question(
+ question_id,
+ question_text,
+ puncts,
+ all_top_chunks,
+ all_chunk_similarities,
+ all_chunk_overlaps,
+ top_chunks_dir
+ )
+
+ # Подсчитываем метрики для текущего вопроса
+ question_found_puncts = sum(punct_found)
+ found_puncts += question_found_puncts
+
+ relevant_chunks += question_relevant_chunks
+
+ # Обновляем метрики для документов
+ question_found_relevant_docs = sum(1 for doc in docs_found_for_question if doc in relevant_docs)
+ found_relevant_docs += question_found_relevant_docs
+ question_total_docs_found = len(docs_found_for_question)
+ total_docs_found += question_total_docs_found
+
+ # Вычисляем метрики для текущего вопроса
+ question_text_precision = question_relevant_chunks / question_total_chunks if question_total_chunks > 0 else 0
+ question_text_recall = question_found_puncts / question_total_puncts if question_total_puncts > 0 else 0
+ question_text_f1 = 2 * question_text_precision * question_text_recall / (question_text_precision + question_text_recall) if question_text_precision + question_text_recall > 0 else 0
+
+ question_doc_precision = question_found_relevant_docs / question_total_docs_found if question_total_docs_found > 0 else 0
+ question_doc_recall = question_found_relevant_docs / question_total_docs_required if question_total_docs_required > 0 else 0
+ question_doc_f1 = 2 * question_doc_precision * question_doc_recall / (question_doc_precision + question_doc_recall) if question_doc_precision + question_doc_recall > 0 else 0
+
+ # Сохраняем метрики вопроса
+ question_metrics.append({
+ 'question_id': question_id,
+ 'question_text': question_text,
+ 'top_n': top_n,
+ 'text_precision': question_text_precision,
+ 'text_recall': question_text_recall,
+ 'text_f1': question_text_f1,
+ 'doc_precision': question_doc_precision,
+ 'doc_recall': question_doc_recall,
+ 'doc_f1': question_doc_f1,
+ 'found_puncts': question_found_puncts,
+ 'total_puncts': question_total_puncts,
+ 'relevant_chunks': question_relevant_chunks,
+ 'total_chunks': question_total_chunks,
+ 'found_relevant_docs': question_found_relevant_docs,
+ 'total_docs_required': question_total_docs_required,
+ 'total_docs_found': question_total_docs_found
+ })
+
+ # Вычисляем метрики для текста
+ text_precision = relevant_chunks / total_chunks if total_chunks > 0 else 0
+ text_recall = found_puncts / total_puncts if total_puncts > 0 else 0
+ text_f1 = 2 * text_precision * text_recall / (text_precision + text_recall) if text_precision + text_recall > 0 else 0
+
+ # Вычисляем метрики для документов
+ doc_precision = found_relevant_docs / total_docs_found if total_docs_found > 0 else 0
+ doc_recall = found_relevant_docs / total_docs_required if total_docs_required > 0 else 0
+ doc_f1 = 2 * doc_precision * doc_recall / (doc_precision + doc_recall) if doc_precision + doc_recall > 0 else 0
+
+ aggregated_metrics = {
+ 'top_n': top_n,
+ 'text_precision': text_precision,
+ 'text_recall': text_recall,
+ 'text_f1': text_f1,
+ 'doc_precision': doc_precision,
+ 'doc_recall': doc_recall,
+ 'doc_f1': doc_f1,
+ 'found_puncts': found_puncts,
+ 'total_puncts': total_puncts,
+ 'relevant_chunks': relevant_chunks,
+ 'total_chunks': total_chunks,
+ 'found_relevant_docs': found_relevant_docs,
+ 'total_docs_required': total_docs_required,
+ 'total_docs_found': total_docs_found
+ }
+
+ return aggregated_metrics, pd.DataFrame(question_metrics)
+
+
+def main():
+ """
+ Основная функция скрипта.
+ """
+ args = parse_args()
+
+ # Устанавливаем устройство из аргументов
+ device = args.device
+
+ # Создаем выходной каталог, если его нет
+ os.makedirs(args.output_dir, exist_ok=True)
+
+ # Создаем директорию для топ-чанков
+ top_chunks_dir = os.path.join(args.output_dir, "top_chunks")
+ os.makedirs(top_chunks_dir, exist_ok=True)
+
+ # Загружаем датасет с вопросами
+ questions_df = load_questions_dataset(args.dataset_path)
+
+ # Формируем уникальное имя для сохраняемых файлов на основе параметров стратегии и модели
+ strategy_config_str = f"fixed_size_w{args.words_per_chunk}_o{args.overlap_words}"
+ chunks_filename = f"chunks_{strategy_config_str}_{args.model_name.replace('/', '_')}"
+ questions_filename = f"questions_{args.model_name.replace('/', '_')}"
+
+ # Пытаемся загрузить сохраненные эмбеддинги и данные
+ chunk_embeddings, chunks_df = None, None
+ question_embeddings, questions_df_with_embeddings = None, None
+
+ if not args.force_recompute:
+ chunk_embeddings, chunks_df = load_embeddings_and_data(chunks_filename, args.output_dir)
+ question_embeddings, questions_df_with_embeddings = load_embeddings_and_data(questions_filename, args.output_dir)
+
+ # Если не удалось загрузить данные или включен режим принудительного пересчета
+ if chunk_embeddings is None or chunks_df is None:
+ # Читаем и обрабатываем документы
+ documents = read_documents(args.data_folder)
+
+ # Формируем конфигурацию для стратегии fixed_size
+ fixed_size_config = {
+ "words_per_chunk": args.words_per_chunk,
+ "overlap_words": args.overlap_words
+ }
+
+ # Получаем DataFrame с чанками
+ chunks_df = process_documents(documents, fixed_size_config)
+
+ # Настраиваем модель и токенизатор
+ model, tokenizer = setup_model_and_tokenizer(args.model_name, args.use_sentence_transformers, device)
+
+ # Получаем эмбеддинги для чанков
+ chunk_embeddings = get_embeddings(chunks_df['text'].tolist(), model, tokenizer, args.batch_size, args.use_sentence_transformers, device)
+
+ # Сохраняем эмбеддинги и данные
+ save_embeddings_and_data(chunk_embeddings, chunks_df, chunks_filename, args.output_dir)
+
+ # Если не удалось загрузить эмбеддинги вопросов или включен режим принудительного пересчета
+ if question_embeddings is None or questions_df_with_embeddings is None:
+ # Получаем уникальные вопросы (по id)
+ unique_questions = questions_df.drop_duplicates(subset=['id'])[['id', 'question']]
+
+ # Настраиваем модель и токенизатор (если еще не настроены)
+ if 'model' not in locals() or 'tokenizer' not in locals():
+ model, tokenizer = setup_model_and_tokenizer(args.model_name, args.use_sentence_transformers, device)
+
+ # Получаем эмбеддинги для вопросов
+ question_embeddings = get_embeddings(unique_questions['question'].tolist(), model, tokenizer, args.batch_size, args.use_sentence_transformers, device)
+
+ # Сохраняем эмбеддинги и данные
+ save_embeddings_and_data(question_embeddings, unique_questions, questions_filename, args.output_dir)
+
+ # Устанавливаем questions_df_with_embeddings для дальнейшего использования
+ questions_df_with_embeddings = unique_questions
+
+ # Создаем словарь соответствия id вопроса и его индекса в эмбеддингах
+ question_id_to_idx = {
+ row['id']: i
+ for i, (_, row) in enumerate(questions_df_with_embeddings.iterrows())
+ }
+
+ # Оцениваем стратегию чанкинга для разных значений top_n
+ aggregated_results = []
+ all_question_metrics = []
+
+ for top_n in TOP_N_VALUES:
+ metrics, question_metrics = evaluate_for_top_n_with_mapping(
+ questions_df, # Исходный датасет с связью между вопросами и документами
+ chunks_df, # Датасет с чанками
+ question_embeddings, # Эмбеддинги вопросов
+ chunk_embeddings, # Эмбеддинги чанков
+ question_id_to_idx, # Маппинг id вопроса к индексу в эмбеддингах
+ top_n, # Количество чанков в топе
+ args.similarity_threshold, # Порог для определения перекрытия
+ top_chunks_dir if top_n == 20 else None # Сохраняем топ-чанки только для top_n=20
+ )
+ aggregated_results.append(metrics)
+ all_question_metrics.append(question_metrics)
+
+ # Объединяем все метрики по вопросам
+ all_question_metrics_df = pd.concat(all_question_metrics)
+
+ # Создаем DataFrame с агрегированными результатами
+ aggregated_results_df = pd.DataFrame(aggregated_results)
+
+ # Сохраняем результаты
+ results_filename = f"results_{strategy_config_str}_{args.model_name.replace('/', '_')}.csv"
+ results_path = os.path.join(args.output_dir, results_filename)
+ aggregated_results_df.to_csv(results_path, index=False)
+
+ # Сохраняем метрики по вопросам
+ question_metrics_filename = f"question_metrics_{strategy_config_str}_{args.model_name.replace('/', '_')}.xlsx"
+ question_metrics_path = os.path.join(args.output_dir, question_metrics_filename)
+ all_question_metrics_df.to_excel(question_metrics_path, index=False)
+
+ print(f"\nРезультаты сохранены в {results_path}")
+ print(f"Метрики по вопросам сохранены в {question_metrics_path}")
+ print(f"Топ-20 чанков для каждого вопроса сохранены в {top_chunks_dir}")
+ print("\nМетрики для различных значений top_n:")
+ print(aggregated_results_df[['top_n', 'text_precision', 'text_recall', 'text_f1', 'doc_precision', 'doc_recall', 'doc_f1']])
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/lib/extractor/scripts/plot_macro_metrics.py b/lib/extractor/scripts/plot_macro_metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..88bb59a4c893d30941dbedc0e3abb99f199311a5
--- /dev/null
+++ b/lib/extractor/scripts/plot_macro_metrics.py
@@ -0,0 +1,348 @@
+#!/usr/bin/env python
+"""
+Скрипт для построения специализированных графиков на основе макрометрик из Excel-файла.
+Строит несколько типов графиков:
+1. Зависимость macro_text_recall от top_N для разных моделей при фиксированных параметрах чанкинга
+2. Зависимость macro_text_recall от top_N для разных подходов к чанкингу при фиксированных моделях
+3. Зависимость macro_text_recall от подхода к чанкингу для разных моделей при фиксированных top_N
+"""
+
+import os
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
+
+# Константы
+EXCEL_FILE_PATH = "../../Белагропромбанк/test_vectors/combined_results.xlsx"
+PLOTS_DIR = "../../Белагропромбанк/test_vectors/plots"
+
+# Настройки для графиков
+plt.rcParams['font.family'] = 'DejaVu Sans'
+sns.set_style("whitegrid")
+FIGSIZE = (14, 10)
+DPI = 300
+
+
+def setup_plots_directory(plots_dir: str) -> None:
+ """
+ Создает директорию для сохранения графиков, если она не существует.
+
+ Args:
+ plots_dir: Путь к директории для графиков
+ """
+ if not os.path.exists(plots_dir):
+ os.makedirs(plots_dir)
+ print(f"Создана директория для графиков: {plots_dir}")
+ else:
+ print(f"Использование существующей директории для графиков: {plots_dir}")
+
+
+def load_macro_metrics(excel_path: str) -> pd.DataFrame:
+ """
+ Загружает макрометрики из Excel-файла.
+
+ Args:
+ excel_path: Путь к Excel-файлу с данными
+
+ Returns:
+ DataFrame с макрометриками
+ """
+ try:
+ df = pd.read_excel(excel_path, sheet_name="Macro метрики")
+ print(f"Загружены данные из {excel_path}, лист 'Macro метрики'")
+ print(f"Количество строк: {len(df)}")
+ return df
+ except Exception as e:
+ print(f"Ошибка при загрузке данных: {e}")
+ raise
+
+
+def plot_top_n_vs_recall_by_model(df: pd.DataFrame, plots_dir: str) -> None:
+ """
+ Строит графики зависимости macro_text_recall от top_N для разных моделей
+ при фиксированных параметрах чанкинга (50/25 и 200/75).
+
+ Args:
+ df: DataFrame с данными
+ plots_dir: Директория для сохранения графиков
+ """
+ # Фиксированные параметры чанкинга
+ chunking_params = [
+ {"words": 50, "overlap": 25, "title": "Чанкинг 50/25"},
+ {"words": 200, "overlap": 75, "title": "Чанкинг 200/75"}
+ ]
+
+ # Создаем субплоты: 1 строка, 2 столбца
+ fig, axes = plt.subplots(1, 2, figsize=FIGSIZE, sharey=True)
+
+ for i, params in enumerate(chunking_params):
+ # Фильтруем данные для текущих параметров чанкинга
+ filtered_df = df[
+ (df['words_per_chunk'] == params['words']) &
+ (df['overlap_words'] == params['overlap'])
+ ]
+
+ if len(filtered_df) == 0:
+ print(f"Предупреждение: нет данных для чанкинга {params['words']}/{params['overlap']}")
+ axes[i].text(0.5, 0.5, f"Нет данных для чанкинга {params['words']}/{params['overlap']}",
+ ha='center', va='center', fontsize=12)
+ axes[i].set_title(params['title'])
+ continue
+
+ # Находим уникальные модели
+ models = filtered_df['model'].unique()
+
+ # Создаем палитру цветов
+ palette = sns.color_palette("viridis", len(models))
+
+ # Строим график для каждой модели
+ for j, model in enumerate(models):
+ model_df = filtered_df[filtered_df['model'] == model].sort_values('top_n')
+
+ if len(model_df) <= 1:
+ print(f"Предупреждение: недостаточно данных для модели {model} при чанкинге {params['words']}/{params['overlap']}")
+ continue
+
+ # Строим ломаную линию
+ axes[i].plot(model_df['top_n'], model_df['macro_text_recall'],
+ marker='o', linestyle='-', linewidth=2,
+ label=model, color=palette[j])
+
+ # Настраиваем оси и заголовок
+ axes[i].set_title(params['title'], fontsize=14)
+ axes[i].set_xlabel('top_N', fontsize=12)
+ if i == 0:
+ axes[i].set_ylabel('macro_text_recall', fontsize=12)
+
+ # Добавляем сетку
+ axes[i].grid(True, linestyle='--', alpha=0.7)
+
+ # Добавляем легенду
+ axes[i].legend(title="Модель", fontsize=10, loc='best')
+
+ # Общий заголовок
+ plt.suptitle('Зависимость macro_text_recall от top_N для разных моделей', fontsize=16)
+
+ # Настраиваем макет
+ plt.tight_layout(rect=[0, 0, 1, 0.96])
+
+ # Сохраняем график
+ file_path = os.path.join(plots_dir, "top_n_vs_recall_by_model.png")
+ plt.savefig(file_path, dpi=DPI)
+ plt.close()
+
+ print(f"Создан график: {file_path}")
+
+
+def plot_top_n_vs_recall_by_chunking(df: pd.DataFrame, plots_dir: str) -> None:
+ """
+ Строит графики зависимости macro_text_recall от top_N для разных параметров чанкинга
+ при фиксированных моделях (bge и frida).
+
+ Args:
+ df: DataFrame с данными
+ plots_dir: Директория для сохранения графиков
+ """
+ # Фиксированные модели
+ models = ["BAAI/bge", "frida"]
+
+ # Создаем субплоты: 1 строка, 2 столбца
+ fig, axes = plt.subplots(1, 2, figsize=FIGSIZE, sharey=True)
+
+ for i, model_name in enumerate(models):
+ # Находим все строки с моделями, содержащими указанное название
+ model_df = df[df['model'].str.contains(model_name, case=False)]
+
+ if len(model_df) == 0:
+ print(f"Предупреждение: нет данных для модели {model_name}")
+ axes[i].text(0.5, 0.5, f"Нет данных для модели {model_name}",
+ ha='center', va='center', fontsize=12)
+ axes[i].set_title(f"Модель: {model_name}")
+ continue
+
+ # Находим уникальные комбинации параметров чанкинга
+ chunking_combinations = model_df.drop_duplicates(['words_per_chunk', 'overlap_words'])[['words_per_chunk', 'overlap_words']]
+
+ # Ограничиваем количество комбинаций до 7 для читаемости
+ if len(chunking_combinations) > 7:
+ print(f"Предупреждение: слишком много комбинаций чанкинга для модели {model_name}, ограничиваем до 7")
+ chunking_combinations = chunking_combinations.head(7)
+
+ # Создаем палитру цветов
+ palette = sns.color_palette("viridis", len(chunking_combinations))
+
+ # Строим график для каждой комбинации параметров чанкинга
+ for j, (_, row) in enumerate(chunking_combinations.iterrows()):
+ words = row['words_per_chunk']
+ overlap = row['overlap_words']
+
+ # Фильтруем данные для текущей модели и параметров чанкинга
+ chunking_df = model_df[
+ (model_df['words_per_chunk'] == words) &
+ (model_df['overlap_words'] == overlap)
+ ].sort_values('top_n')
+
+ if len(chunking_df) <= 1:
+ print(f"Предупреждение: недостаточно данных для модели {model_name} с чанкингом {words}/{overlap}")
+ continue
+
+ # Строим ломаную линию
+ axes[i].plot(chunking_df['top_n'], chunking_df['macro_text_recall'],
+ marker='o', linestyle='-', linewidth=2,
+ label=f"w={words}, o={overlap}", color=palette[j])
+
+ # Настраиваем оси и заголовок
+ axes[i].set_title(f"Модель: {model_name}", fontsize=14)
+ axes[i].set_xlabel('top_N', fontsize=12)
+ if i == 0:
+ axes[i].set_ylabel('macro_text_recall', fontsize=12)
+
+ # Добавляем сетку
+ axes[i].grid(True, linestyle='--', alpha=0.7)
+
+ # Добавляем легенду
+ axes[i].legend(title="Чанкинг", fontsize=10, loc='best')
+
+ # Общий заголовок
+ plt.suptitle('Зависимость macro_text_recall от top_N для разных параметров чанкинга', fontsize=16)
+
+ # Настраиваем макет
+ plt.tight_layout(rect=[0, 0, 1, 0.96])
+
+ # Сохраняем график
+ file_path = os.path.join(plots_dir, "top_n_vs_recall_by_chunking.png")
+ plt.savefig(file_path, dpi=DPI)
+ plt.close()
+
+ print(f"Создан график: {file_path}")
+
+
+def plot_chunking_vs_recall_by_model(df: pd.DataFrame, plots_dir: str) -> None:
+ """
+ Строит графики зависимости macro_text_recall от подхода к чанкингу
+ для разных моделей при фиксированных top_N (5, 20, 100).
+
+ Args:
+ df: DataFrame с данными
+ plots_dir: Директория для сохранения графиков
+ """
+ # Фиксированные значения top_N
+ top_n_values = [5, 20, 100]
+
+ # Создаем субплоты: 1 строка, 3 столбца
+ fig, axes = plt.subplots(1, 3, figsize=FIGSIZE, sharey=True)
+
+ # Создаем порядок чанкинга - сортируем по возрастанию размера и оверлапа
+ chunking_order = df.drop_duplicates(['words_per_chunk', 'overlap_words'])[['words_per_chunk', 'overlap_words']]
+ chunking_order = chunking_order.sort_values(['words_per_chunk', 'overlap_words'])
+
+ # Создаем словарь для маппинга комбинаций чанкинга на индексы
+ chunking_labels = [f"{row['words_per_chunk']}/{row['overlap_words']}" for _, row in chunking_order.iterrows()]
+ chunking_map = {f"{row['words_per_chunk']}/{row['overlap_words']}": i for i, (_, row) in enumerate(chunking_order.iterrows())}
+
+ for i, top_n in enumerate(top_n_values):
+ # Фильтруем данные для текущего top_N
+ top_n_df = df[df['top_n'] == top_n]
+
+ if len(top_n_df) == 0:
+ print(f"Предупреждение: нет данных для top_N={top_n}")
+ axes[i].text(0.5, 0.5, f"Нет данных для top_N={top_n}",
+ ha='center', va='center', fontsize=12)
+ axes[i].set_title(f"top_N={top_n}")
+ continue
+
+ # Находим уникальные модели
+ models = top_n_df['model'].unique()
+
+ # Ограничиваем количество моделей до 5 для читаемости
+ if len(models) > 5:
+ print(f"Предупреждение: слишком много моделей для top_N={top_n}, ограничиваем до 5")
+ models = models[:5]
+
+ # Создаем палитру цветов
+ palette = sns.color_palette("viridis", len(models))
+
+ # Строим график для каждой модели
+ for j, model in enumerate(models):
+ model_df = top_n_df[top_n_df['model'] == model].copy()
+
+ if len(model_df) <= 1:
+ print(f"Предупреждение: недостаточно данных для модели {model} при top_N={top_n}")
+ continue
+
+ # Создаем новую колонку с индексом чанкинга для сортировки
+ model_df['chunking_index'] = model_df.apply(
+ lambda row: chunking_map.get(f"{row['words_per_chunk']}/{row['overlap_words']}", -1),
+ axis=1
+ )
+
+ # Отбрасываем строки с неизвестными комбинациями чанкинга
+ model_df = model_df[model_df['chunking_index'] >= 0]
+
+ if len(model_df) <= 1:
+ continue
+
+ # Сортируем по индексу чанкинга
+ model_df = model_df.sort_values('chunking_index')
+
+ # Создаем список индексов и значений для графика
+ x_indices = model_df['chunking_index'].tolist()
+ y_values = model_df['macro_text_recall'].tolist()
+
+ # Строим ломаную линию
+ axes[i].plot(x_indices, y_values, marker='o', linestyle='-', linewidth=2,
+ label=model, color=palette[j])
+
+ # Настраиваем оси и заголовок
+ axes[i].set_title(f"top_N={top_n}", fontsize=14)
+ axes[i].set_xlabel('Подход к чанкингу', fontsize=12)
+ if i == 0:
+ axes[i].set_ylabel('macro_text_recall', fontsize=12)
+
+ # Устанавливаем метки на оси X (подходы к чанкингу)
+ axes[i].set_xticks(range(len(chunking_labels)))
+ axes[i].set_xticklabels(chunking_labels, rotation=45, ha='right', fontsize=10)
+
+ # Добавляем сетку
+ axes[i].grid(True, linestyle='--', alpha=0.7)
+
+ # Добавляем легенду
+ axes[i].legend(title="Модель", fontsize=10, loc='best')
+
+ # Общий заголовок
+ plt.suptitle('Зависимость macro_text_recall от подхода к чанкингу для разных моделей', fontsize=16)
+
+ # Настраиваем макет
+ plt.tight_layout(rect=[0, 0, 1, 0.96])
+
+ # Сохраняем график
+ file_path = os.path.join(plots_dir, "chunking_vs_recall_by_model.png")
+ plt.savefig(file_path, dpi=DPI)
+ plt.close()
+
+ print(f"Создан график: {file_path}")
+
+
+def main():
+ """Основная функция скрипта."""
+ # Создаем директорию для графиков
+ setup_plots_directory(PLOTS_DIR)
+
+ # Загружаем данные
+ try:
+ macro_metrics = load_macro_metrics(EXCEL_FILE_PATH)
+ except Exception as e:
+ print(f"Критическая ошибка: {e}")
+ return
+
+ # Строим графики
+ plot_top_n_vs_recall_by_model(macro_metrics, PLOTS_DIR)
+ plot_top_n_vs_recall_by_chunking(macro_metrics, PLOTS_DIR)
+ plot_chunking_vs_recall_by_model(macro_metrics, PLOTS_DIR)
+
+ print("Готово! Все графики созданы.")
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/lib/extractor/scripts/prepare_dataset.py b/lib/extractor/scripts/prepare_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..af739425cea83138ecb68119547b717afa04f985
--- /dev/null
+++ b/lib/extractor/scripts/prepare_dataset.py
@@ -0,0 +1,578 @@
+#!/usr/bin/env python
+"""
+Скрипт для подготовки датасета с вопросами и текстами пунктов/приложений.
+Преобразует исходный датасет, содержащий списки пунктов, в расширенный датасет,
+где каждому пункту/приложению соответствует отдельная строка.
+"""
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Any, Dict
+
+import pandas as pd
+from tqdm import tqdm
+
+from ntr_text_fragmentation import Destructurer
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from ntr_fileparser import UniversalParser
+
+
+def parse_args():
+ """
+ Парсит аргументы командной строки.
+
+ Returns:
+ Аргументы командной строки
+ """
+ parser = argparse.ArgumentParser(description="Подготовка датасета с текстами пунктов")
+
+ parser.add_argument('--input-dataset', type=str, default='data/dataset.xlsx',
+ help='Путь к исходному датасету (Excel-файл)')
+ parser.add_argument('--output-dataset', type=str, default='data/dataset_with_texts.xlsx',
+ help='Путь для сохранения подготовленного датасета (Excel-файл)')
+ parser.add_argument('--data-folder', type=str, default='data/docs',
+ help='Путь к папке с документами')
+ parser.add_argument('--debug', action='store_true',
+ help='Включить режим отладки с дополнительным выводом информации')
+
+ return parser.parse_args()
+
+
+def load_dataset(file_path: str, debug: bool = False) -> pd.DataFrame:
+ """
+ Загружает исходный датасет с вопросами.
+
+ Args:
+ file_path: Путь к Excel-файлу
+ debug: Режим отладки
+
+ Returns:
+ DataFrame с вопросами
+ """
+ print(f"Загрузка исходного датасета из {file_path}...")
+
+ df = pd.read_excel(file_path)
+
+ # Преобразуем строковые списки в настоящие списки
+ for col in ['puncts', 'appendices']:
+ if col in df.columns:
+ df[col] = df[col].apply(lambda x:
+ eval(x) if isinstance(x, str) and x.strip()
+ else ([] if pd.isna(x) else x))
+
+ # Вывод отладочной информации о форматах пунктов/приложений
+ if debug:
+ all_puncts = set()
+ all_appendices = set()
+
+ for _, row in df.iterrows():
+ if 'puncts' in row and row['puncts']:
+ all_puncts.update(row['puncts'])
+ if 'appendices' in row and row['appendices']:
+ all_appendices.update(row['appendices'])
+
+ print(f"\nУникальные форматы пунктов в датасете ({len(all_puncts)}):")
+ for i, p in enumerate(sorted(all_puncts)):
+ if i < 20 or i > len(all_puncts) - 20:
+ print(f" - {repr(p)}")
+ elif i == 20:
+ print(" ... (пропущено)")
+
+ print(f"\nУникальные форматы приложений в датасете ({len(all_appendices)}):")
+ for app in sorted(all_appendices):
+ print(f" - {repr(app)}")
+
+ print(f"Загружено {len(df)} вопросов")
+ return df
+
+
+def read_documents(folder_path: str) -> Dict[str, Any]:
+ """
+ Читает все документы из указанной папки.
+
+ Args:
+ folder_path: Путь к папке с документами
+
+ Returns:
+ Словарь {имя_файла: parsed_document}
+ """
+ print(f"Чтение документов из {folder_path}...")
+ parser = UniversalParser()
+ documents = {}
+
+ for file_path in tqdm(list(Path(folder_path).glob("*.docx")), desc="Чтение документов"):
+ try:
+ doc_name = file_path.stem
+ documents[doc_name] = parser.parse_by_path(str(file_path))
+ except Exception as e:
+ print(f"Ошибка при чтении файла {file_path}: {e}")
+
+ print(f"Прочитано {len(documents)} документов")
+ return documents
+
+
+def normalize_punct_format(punct: str) -> str:
+ """
+ Нормализует формат номера пункта для единообразного сравнения.
+
+ Args:
+ punct: Номер пункта
+
+ Returns:
+ Нормализованный номер пункта
+ """
+ # Убираем пробелы
+ punct = punct.strip()
+
+ # Убираем завершающую точку, если она есть
+ if punct.endswith('.'):
+ punct = punct[:-1]
+
+ return punct
+
+
+def normalize_appendix_format(appendix: str) -> str:
+ """
+ Нормализует формат номера приложения для единообразного сравнения.
+
+ Args:
+ appendix: Номер приложения
+
+ Returns:
+ Нормализованный номер приложения
+ """
+ # Убираем пробелы
+ appendix = appendix.strip()
+
+ # Обработка форматов с дефисом (например, "14-1")
+ if "-" in appendix:
+ return appendix
+
+ return appendix
+
+
+def find_matching_key(search_key, available_keys, item_type='punct', debug_mode=False):
+ """
+ Ищет наиболее подходящий ключ среди доступных ключей с учетом типа элемента
+
+ Args:
+ search_key: Ключ для поиска
+ available_keys: Доступные ключи
+ item_type: Тип элемента ('punct' или 'appendix')
+ debug_mode: Режим отладки
+
+ Returns:
+ Найденный ключ или None
+ """
+ if not available_keys:
+ return None
+
+ # Нормализуем ключ в зависимости от типа элемента
+ if item_type == 'punct':
+ normalized_search_key = normalize_punct_format(search_key)
+ else: # appendix
+ normalized_search_key = normalize_appendix_format(search_key)
+
+ # Проверяем прямое совпадение ключей
+ for key in available_keys:
+ if item_type == 'punct':
+ normalized_key = normalize_punct_format(key)
+ else: # appendix
+ normalized_key = normalize_appendix_format(key)
+
+ if normalized_key == normalized_search_key:
+ if debug_mode:
+ print(f"Найдено прямое совпадение для {item_type} {search_key} -> {key}")
+ return key
+
+ # Если прямого совпадения нет, проверяем "мягкое" совпадение
+ # Только для пунктов, не для приложений
+ if item_type == 'punct':
+ for key in available_keys:
+ normalized_key = normalize_punct_format(key)
+
+ # Если ключ содержит "/", это подпункт приложения, его не следует сопоставлять с обычным пунктом
+ if '/' in key and '/' not in search_key:
+ continue
+
+ # Проверяем совпадение конца номера (например, "1.2" и "1.2.")
+ if normalized_key.rstrip('.') == normalized_search_key.rstrip('.'):
+ if debug_mode:
+ print(f"Найдено мягкое совпадение для {search_key} -> {key}")
+ return key
+
+ return None
+
+
+def extract_item_texts(documents, debug_mode=False):
+ """
+ Извлекает тексты пунктов и приложений из документов.
+
+ Args:
+ documents: Словарь с распарсенными документами {doc_name: document}
+ debug_mode: Включать ли режим отладки
+
+ Returns:
+ Словарь с текстами пунктов и приложений, организованный по названиям документов
+ """
+ print("Извлечение текстов пунктов и приложений...")
+
+ item_texts = {}
+ all_extracted_items = set()
+ all_extracted_appendices = set()
+
+ for doc_name, document in tqdm(documents.items(), desc="Применение стратегии numbered_items"):
+ # Используем стратегию numbered_items с режимом отладки
+ destructurer = Destructurer(document)
+ destructurer.configure('numbered_items', debug_mode=debug_mode)
+ entities, _ = destructurer.destructure()
+
+ # Инициализируем структуру для документа, если она еще не создана
+ if doc_name not in item_texts:
+ item_texts[doc_name] = {
+ 'puncts': {}, # Для пунктов основного текста
+ 'appendices': {} # Для приложений
+ }
+
+ for entity in entities:
+ # Пропускаем сущность документа
+ if entity.type == "Document":
+ continue
+
+ # Работаем только с чанками для поиска
+ if hasattr(entity, 'use_in_search') and entity.use_in_search:
+ metadata = entity.metadata
+ text = entity.text
+
+ # Для пунктов
+ if 'item_number' in metadata:
+ item_number = metadata['item_number']
+
+ # Проверяем, является ли пункт подпунктом приложения
+ if 'appendix_number' in metadata:
+ # Это подпункт приложения
+ appendix_number = metadata['appendix_number']
+
+ # Создаем структуру для приложения, если ее еще нет
+ if appendix_number not in item_texts[doc_name]['appendices']:
+ item_texts[doc_name]['appendices'][appendix_number] = {
+ 'main_text': '', # Основной текст приложения
+ 'subpuncts': {} # Подпункты приложения
+ }
+
+ # Добавляем подпункт в словарь подпунктов
+ item_texts[doc_name]['appendices'][appendix_number]['subpuncts'][item_number] = text
+
+ if debug_mode:
+ print(f"Извлечен подпункт {item_number} приложения {appendix_number} из {doc_name}")
+
+ all_extracted_items.add(item_number)
+ else:
+ # Обычный пункт
+ item_texts[doc_name]['puncts'][item_number] = text
+
+ if debug_mode:
+ print(f"Извлечен пункт {item_number} из {doc_name}")
+
+ all_extracted_items.add(item_number)
+
+ # Для приложений
+ elif 'appendix_number' in metadata and 'item_number' not in metadata:
+ appendix_number = metadata['appendix_number']
+
+ # Создаем структуру для приложения, если ее еще нет
+ if appendix_number not in item_texts[doc_name]['appendices']:
+ item_texts[doc_name]['appendices'][appendix_number] = {
+ 'main_text': text, # Основной текст приложения
+ 'subpuncts': {} # Подпункты приложения
+ }
+ else:
+ # Если приложение уже существует, обновляем основной текст
+ item_texts[doc_name]['appendices'][appendix_number]['main_text'] = text
+
+ if debug_mode:
+ print(f"Извлечено приложение {appendix_number} из {doc_name}")
+
+ all_extracted_appendices.add(appendix_number)
+
+ # Выводим статистику, если включен режим отладки
+ if debug_mode:
+ print(f"\nВсего извлечено уникальных пунктов: {len(all_extracted_items)}")
+ print(f"Примеры форматов пунктов: {', '.join(sorted(list(all_extracted_items))[:20])}")
+
+ print(f"\nВсего извлечено уникальных приложений: {len(all_extracted_appendices)}")
+ print(f"Форматы приложений: {', '.join(sorted(list(all_extracted_appendices)))}")
+
+ # Подсчитываем общее количество пунктов и приложений
+ total_puncts = sum(len(doc_data['puncts']) for doc_data in item_texts.values())
+ total_appendices = sum(len(doc_data['appendices']) for doc_data in item_texts.values())
+
+ print(f"Извлечено {total_puncts} пунктов и {total_appendices} приложений из {len(item_texts)} документов")
+
+ return item_texts
+
+
+def is_subpunct(parent_punct: str, possible_subpunct: str) -> bool:
+ """
+ Проверяет, является ли пункт подпунктом другого пункта.
+
+ Args:
+ parent_punct: Родительский пункт (например, "14")
+ possible_subpunct: Возможный подпункт (например, "14.1")
+
+ Returns:
+ True, если possible_subpunct является подпунктом parent_punct
+ """
+ # Нормализуем пункты
+ parent = normalize_punct_format(parent_punct)
+ child = normalize_punct_format(possible_subpunct)
+
+ # Проверяем, начинается ли child с parent и после него идет точка или другой разделитель
+ if child.startswith(parent):
+ # Если длины равны, это тот же самый пункт
+ if len(child) == len(parent):
+ return False
+
+ # Проверяем символ после parent - должна быть точка (дефис исключен, т.к. это разные пункты)
+ next_char = child[len(parent)]
+ return next_char in ['.']
+
+ return False
+
+
+def collect_subpuncts(punct: str, all_puncts: dict) -> dict:
+ """
+ Собирает все подпункты для указанного пункта.
+
+ Args:
+ punct: Пункт, для которого нужно найти подпункты (например, "14")
+ all_puncts: Словарь всех пунктов {punct: text}
+
+ Returns:
+ Словарь {punct: text} с пунктом и всеми его подпунктами
+ """
+ result = {}
+ normalized_punct = normalize_punct_format(punct)
+
+ # Добавляем сам пункт, если он существует
+ if normalized_punct in all_puncts:
+ result[normalized_punct] = all_puncts[normalized_punct]
+
+ # Ищем подпункты
+ for possible_subpunct in all_puncts.keys():
+ if is_subpunct(normalized_punct, possible_subpunct):
+ result[possible_subpunct] = all_puncts[possible_subpunct]
+
+ return result
+
+
+def prepare_expanded_dataset(df, item_texts, output_path, debug_mode=False):
+ """
+ Подготавливает расширенный датасет, добавляя тексты пунктов и приложений.
+
+ Args:
+ df: Исходный датасет
+ item_texts: Словарь с текстами пунктов и приложений
+ output_path: Путь для сохранения расширенного датасета
+ debug_mode: Включать ли режим отладки
+
+ Returns:
+ Датафрейм с расширенным датасетом
+ """
+ rows = []
+ skipped_items = 0
+ total_items = 0
+
+ for _, row in df.iterrows():
+ question_id = row['id']
+ question = row['question']
+ filepath = row.get('filepath', '')
+
+ # Получаем имя файла без пути
+ doc_name = Path(filepath).stem if filepath else ''
+
+ # Пропускаем, если файл не найден
+ if not doc_name or doc_name not in item_texts:
+ if debug_mode and doc_name:
+ print(f"Документ {doc_name} не найден в извлеченных данных")
+ continue
+
+ # Обрабатываем пункты
+ puncts = row.get('puncts', [])
+ if isinstance(puncts, str) and puncts.strip():
+ # Преобразуем строковое представление в список
+ try:
+ puncts = eval(puncts)
+ except:
+ puncts = []
+
+ if not isinstance(puncts, list):
+ puncts = []
+
+ for punct in puncts:
+ total_items += 1
+
+ if debug_mode:
+ print(f"\nОбработка пункта {punct} для вопроса {question_id} из {doc_name}")
+
+ # Ищем соответствующий пункт в документе
+ available_keys = list(item_texts[doc_name]['puncts'].keys())
+ matching_key = find_matching_key(punct, available_keys, 'punct', debug_mode)
+
+ if matching_key:
+ # Сохраняем основной текст пункта
+ item_text = item_texts[doc_name]['puncts'][matching_key]
+
+ # Список всех включенных ключей (для отслеживания что было приконкатенировано)
+ matched_keys = [matching_key]
+
+ # Ищем все подпункты для этого пункта
+ subpuncts = {}
+ for key in available_keys:
+ if is_subpunct(matching_key, key):
+ subpuncts[key] = item_texts[doc_name]['puncts'][key]
+ matched_keys.append(key)
+
+ # Если есть подпункты, добавляем их к основному тексту
+ if subpuncts:
+ # Сортируем подпункты по номеру
+ sorted_subpuncts = sorted(subpuncts.items(), key=lambda x: x[0])
+
+ # Добавляем разделитель и все подпункты
+ combined_text = item_text
+ for key, subtext in sorted_subpuncts:
+ combined_text += f"\n\n{key} {subtext}"
+
+ item_text = combined_text
+
+ # Добавляем строку с пунктом и его подпунктами
+ rows.append({
+ 'id': question_id,
+ 'question': question,
+ 'filename': doc_name,
+ 'text': item_text,
+ 'item_type': 'punct',
+ 'item_id': punct,
+ 'matching_keys': ", ".join(matched_keys)
+ })
+
+ if debug_mode:
+ print(f"Добавлен пункт {matching_key} для {question_id} с {len(matched_keys)} ключами")
+ if len(matched_keys) > 1:
+ print(f" Включены ключи: {', '.join(matched_keys)}")
+ else:
+ skipped_items += 1
+ if debug_mode:
+ print(f"Не найден соответствующий пункт для {punct} в {doc_name}")
+
+ # Обрабатываем приложения
+ appendices = row.get('appendices', [])
+ if isinstance(appendices, str) and appendices.strip():
+ # Преобразуем строковое представление в список
+ try:
+ appendices = eval(appendices)
+ except:
+ appendices = []
+
+ if not isinstance(appendices, list):
+ appendices = []
+
+ for appendix in appendices:
+ total_items += 1
+
+ if debug_mode:
+ print(f"\nОбработка приложения {appendix} для вопроса {question_id} из {doc_name}")
+
+ # Ищем соответствующее приложение в документе
+ available_keys = list(item_texts[doc_name]['appendices'].keys())
+ matching_key = find_matching_key(appendix, available_keys, 'appendix', debug_mode)
+
+ if matching_key:
+ appendix_content = item_texts[doc_name]['appendices'][matching_key]
+
+ # Список всех включенных ключей (для отслеживания что было приконкатенировано)
+ matched_keys = [matching_key]
+
+ # Формируем полный текст приложения, включая все подпункты
+ if isinstance(appendix_content, dict):
+ # Начинаем с основного текста
+ full_text = appendix_content.get('main_text', '')
+
+ # Добавляем все подпункты в отсортированном порядке
+ if 'subpuncts' in appendix_content and appendix_content['subpuncts']:
+ subpuncts = appendix_content['subpuncts']
+ sorted_subpuncts = sorted(subpuncts.items(), key=lambda x: x[0])
+
+ # Добавляем разделитель, если есть основной текст
+ if full_text:
+ full_text += "\n\n"
+
+ # Добавляем все подпункты
+ for i, (key, subtext) in enumerate(sorted_subpuncts):
+ matched_keys.append(f"{matching_key}/{key}")
+ if i > 0:
+ full_text += "\n\n"
+ full_text += f"{key} {subtext}"
+ else:
+ # Если приложение просто строка
+ full_text = appendix_content
+
+ # Добавляем строку с приложением
+ rows.append({
+ 'id': question_id,
+ 'question': question,
+ 'filename': doc_name,
+ 'text': full_text,
+ 'item_type': 'appendix',
+ 'item_id': appendix,
+ 'matching_keys': ", ".join(matched_keys)
+ })
+
+ if debug_mode:
+ print(f"Добавлено приложение {matching_key} для {question_id} с {len(matched_keys)} ключами")
+ if len(matched_keys) > 1:
+ print(f" Включены ключи: {', '.join(matched_keys)}")
+ else:
+ skipped_items += 1
+ if debug_mode:
+ print(f"Не найдено соответствующее приложение для {appendix} в {doc_name}")
+
+ extended_df = pd.DataFrame(rows)
+
+ # Сохраняем расширенный датасет
+ extended_df.to_excel(output_path, index=False)
+
+ print(f"Расширенный датасет сохранен в {output_path}")
+ print(f"Всего обработано элементов: {total_items}")
+ print(f"Всего элементов в расширенном датасете: {len(extended_df)}")
+ print(f"Пропущено элементов из-за отсутствия соответствия: {skipped_items}")
+
+ return extended_df
+
+
+def main():
+ # Парсим аргументы командной строки
+ args = parse_args()
+
+ # Определяем режим отладки
+ debug = args.debug
+
+ # Загружаем исходный датасет
+ df = load_dataset(args.input_dataset, debug)
+
+ # Читаем документы
+ documents = read_documents(args.data_folder)
+
+ # Извлекаем тексты пунктов и приложений
+ item_texts = extract_item_texts(documents, debug)
+
+ # Подготавливаем расширенный датасет
+ expanded_df = prepare_expanded_dataset(df, item_texts, args.output_dataset, debug)
+
+ print("Готово!")
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/lib/extractor/scripts/run_chunking_experiments.sh b/lib/extractor/scripts/run_chunking_experiments.sh
new file mode 100644
index 0000000000000000000000000000000000000000..0f3d9ce4deb7ae6e5639f5f39d93485f92ef4306
--- /dev/null
+++ b/lib/extractor/scripts/run_chunking_experiments.sh
@@ -0,0 +1,156 @@
+#!/bin/bash
+# Скрипт для запуска экспериментов по оценке качества чанкинга с разными моделями и параметрами
+
+# Директории и пути по умолчанию
+DATA_FOLDER="data/docs"
+DATASET_PATH="data/dataset.xlsx"
+OUTPUT_DIR="data"
+LOG_DIR="logs"
+SIMILARITY_THRESHOLD=0.7
+DEVICE="cuda:1"
+
+# Создаем директории, если они не существуют
+mkdir -p "$OUTPUT_DIR"
+mkdir -p "$LOG_DIR"
+
+# Список моделей для тестирования
+MODELS=(
+ "intfloat/e5-base"
+ "intfloat/e5-large"
+ "BAAI/bge-m3"
+ "deepvk/USER-bge-m3"
+ "ai-forever/FRIDA"
+)
+
+# Параметры чанкинга (отсортированы в запрошенном порядке)
+# Формат: [слов_в_чанке]:[нахлест]:[описание]
+CHUNKING_PARAMS=(
+ "50:25:Маленький чанкинг с нахлёстом 50%"
+ "50:0:Маленький чанкинг без нахлёста"
+ "20:10:Очень мелкий чанкинг с нахлёстом 50%"
+ "100:0:Средний чанкинг без нахлёста"
+ "100:25:Средний чанкинг с нахлёстом 25%"
+ "150:50:Крупный чанкинг с нахлёстом 33%"
+ "200:75:Очень крупный чанкинг с нахлёстом 37.5%"
+)
+
+# Функция для запуска одного эксперимента
+run_experiment() {
+ local model="$1"
+ local words="$2"
+ local overlap="$3"
+ local description="$4"
+
+ # Заменяем слеши в имени модели на подчеркивания для имен файлов
+ local model_safe_name=$(echo "$model" | tr '/' '_')
+
+ # Формируем имя файла результатов
+ local results_filename="results_fixed_size_w${words}_o${overlap}_${model_safe_name}.csv"
+ local results_path="${OUTPUT_DIR}/${results_filename}"
+
+ # Формируем имя файла лога
+ local timestamp=$(date +"%Y%m%d_%H%M%S")
+ local log_filename="log_${model_safe_name}_w${words}_o${overlap}_${timestamp}.txt"
+ local log_path="${LOG_DIR}/${log_filename}"
+
+ echo "=============================================================================="
+ echo "Запуск эксперимента:"
+ echo " Модель: $model"
+ echo " Чанкинг: $description (words=$words, overlap=$overlap)"
+ echo " Устройство: $DEVICE"
+ echo " Результаты будут сохранены в: $results_path"
+ echo " Лог: $log_path"
+ echo "=============================================================================="
+
+ # Базовая команда запуска
+ local cmd="python scripts/evaluate_chunking.py \
+ --data-folder \"$DATA_FOLDER\" \
+ --model-name \"$model\" \
+ --dataset-path \"$DATASET_PATH\" \
+ --output-dir \"$OUTPUT_DIR\" \
+ --words-per-chunk $words \
+ --overlap-words $overlap \
+ --similarity-threshold $SIMILARITY_THRESHOLD \
+ --device $DEVICE \
+ --force-recompute"
+
+ # Специальная обработка для модели ai-forever/FRIDA
+ if [[ "$model" == "ai-forever/FRIDA" ]]; then
+ cmd="$cmd --use-sentence-transformers"
+ fi
+
+ # Записываем информацию о запуске в лог
+ echo "Эксперимент запущен в: $(date)" > "$log_path"
+ echo "Команда: $cmd" >> "$log_path"
+ echo "" >> "$log_path"
+
+ # Записываем время начала
+ start_time=$(date +%s)
+
+ # Запускаем команду и записываем вывод в лог
+ eval "$cmd" 2>&1 | tee -a "$log_path"
+ exit_code=${PIPESTATUS[0]}
+
+ # Записываем время окончания
+ end_time=$(date +%s)
+ duration=$((end_time - start_time))
+ duration_min=$(echo "scale=2; $duration/60" | bc)
+
+ # Добавляем информацию о завершении в лог
+ echo "" >> "$log_path"
+ echo "Эксперимент завершен в: $(date)" >> "$log_path"
+ echo "Длительность: $duration секунд ($duration_min минут)" >> "$log_path"
+ echo "Код возврата: $exit_code" >> "$log_path"
+
+ if [ $exit_code -eq 0 ]; then
+ echo "Эксперимент успешно завершен за $duration_min минут"
+ else
+ echo "Эксперимент завершился с ошибкой (код $exit_code)"
+ fi
+}
+
+# Основная функция
+main() {
+ local total_experiments=$((${#MODELS[@]} * ${#CHUNKING_PARAMS[@]}))
+ local completed_experiments=0
+
+ echo "Запуск $total_experiments экспериментов..."
+
+ # Засекаем время начала всех экспериментов
+ local start_time_all=$(date +%s)
+
+ # Сначала перебираем все параметры чанкинга
+ for chunking_param in "${CHUNKING_PARAMS[@]}"; do
+ # Разбиваем строку параметров на составляющие
+ IFS=':' read -r words overlap description <<< "$chunking_param"
+
+ echo -e "\n=== Стратегия чанкинга: $description (words=$words, overlap=$overlap) ===\n"
+
+ # Затем перебираем все модели для текущей стратегии чанкинга
+ for model in "${MODELS[@]}"; do
+ # Запускаем эксперимент
+ run_experiment "$model" "$words" "$overlap" "$description"
+
+ # Увеличиваем счетчик завершенных экспериментов
+ completed_experiments=$((completed_experiments + 1))
+ remaining_experiments=$((total_experiments - completed_experiments))
+
+ if [ $remaining_experiments -gt 0 ]; then
+ echo "Завершено $completed_experiments/$total_experiments экспериментов. Осталось: $remaining_experiments"
+ fi
+ done
+ done
+
+ # Рассчитываем общее время выполнения
+ local end_time_all=$(date +%s)
+ local total_duration=$((end_time_all - start_time_all))
+ local total_duration_min=$(echo "scale=2; $total_duration/60" | bc)
+
+ echo ""
+ echo "Все эксперименты завершены за $total_duration_min минут"
+ echo "Результаты сохранены в $OUTPUT_DIR"
+ echo "Логи сохранены в $LOG_DIR"
+}
+
+# Запускаем основную функцию
+main
\ No newline at end of file
diff --git a/lib/extractor/scripts/run_experiments.py b/lib/extractor/scripts/run_experiments.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ecc946c150cb5a0a9d7ff51df6e899a55678518
--- /dev/null
+++ b/lib/extractor/scripts/run_experiments.py
@@ -0,0 +1,206 @@
+#!/usr/bin/env python
+"""
+Скрипт для запуска экспериментов по оценке качества чанкинга с разными моделями и параметрами.
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+import time
+from datetime import datetime
+
+# Конфигурация моделей
+MODELS = [
+ "intfloat/e5-base",
+ "intfloat/e5-large",
+ "BAAI/bge-m3",
+ "deepvk/USER-bge-m3",
+ "ai-forever/FRIDA"
+]
+
+# Параметры чанкинга (отсортированы в запрошенном порядке)
+CHUNKING_PARAMS = [
+ {"words": 50, "overlap": 25, "description": "Маленький чанкинг с нахлёстом 50%"},
+ {"words": 50, "overlap": 0, "description": "Маленький чанкинг без нахлёста"},
+ {"words": 20, "overlap": 10, "description": "Очень мелкий чанкинг с нахлёстом 50%"},
+ {"words": 100, "overlap": 0, "description": "Средний чанкинг без нахлёста"},
+ {"words": 100, "overlap": 25, "description": "Средний чанкинг с нахлёстом 25%"},
+ {"words": 150, "overlap": 50, "description": "Крупный чанкинг с нахлёстом 33%"},
+ {"words": 200, "overlap": 75, "description": "Очень крупный чанкинг с нахлёстом 37.5%"}
+]
+
+# Значение порога для нечеткого сравнения
+SIMILARITY_THRESHOLD = 0.7
+
+
+def parse_args():
+ """Парсит аргументы командной строки."""
+ parser = argparse.ArgumentParser(description="Запуск экспериментов для оценки качества чанкинга")
+
+ parser.add_argument("--data-folder", type=str, default="data/docs",
+ help="Путь к папке с документами (по умолчанию: data/docs)")
+ parser.add_argument("--dataset-path", type=str, default="data/dataset.xlsx",
+ help="Путь к Excel-датасету с вопросами (по умолчанию: data/dataset.xlsx)")
+ parser.add_argument("--output-dir", type=str, default="data",
+ help="Директория для сохранения результатов (по умолчанию: data)")
+ parser.add_argument("--log-dir", type=str, default="logs",
+ help="Директория для сохранения логов (по умолчанию: logs)")
+ parser.add_argument("--skip-existing", action="store_true",
+ help="Пропускать эксперименты, если файлы результатов уже существуют")
+ parser.add_argument("--similarity-threshold", type=float, default=SIMILARITY_THRESHOLD,
+ help=f"Порог для нечеткого сравнения (по умолчанию: {SIMILARITY_THRESHOLD})")
+ parser.add_argument("--model", type=str, default=None,
+ help="Запустить эксперимент только для указанной модели")
+ parser.add_argument("--chunking-index", type=int, default=None,
+ help="Запустить эксперимент только для указанного индекса конфигурации чанкинга (0-6)")
+ parser.add_argument("--device", type=str, default="cuda:1",
+ help="Устройство для вычислений (по умолчанию: cuda:1)")
+
+ return parser.parse_args()
+
+
+def run_experiment(model_name, chunking_params, args):
+ """
+ Запускает эксперимент с определенной моделью и параметрами чанкинга.
+
+ Args:
+ model_name: Название модели
+ chunking_params: Словарь с параметрами чанкинга
+ args: Аргументы командной строки
+ """
+ words = chunking_params["words"]
+ overlap = chunking_params["overlap"]
+ description = chunking_params["description"]
+
+ # Формируем имя файла результатов
+ results_filename = f"results_fixed_size_w{words}_o{overlap}_{model_name.replace('/', '_')}.csv"
+ results_path = os.path.join(args.output_dir, results_filename)
+
+ # Проверяем, существует ли файл результатов
+ if args.skip_existing and os.path.exists(results_path):
+ print(f"Пропуск: {results_path} уже существует")
+ return
+
+ # Создаем директорию для логов, если она не существует
+ os.makedirs(args.log_dir, exist_ok=True)
+
+ # Формируем имя файла лога
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+ log_filename = f"log_{model_name.replace('/', '_')}_w{words}_o{overlap}_{timestamp}.txt"
+ log_path = os.path.join(args.log_dir, log_filename)
+
+ # Используем тот же интерпретатор Python, что и текущий скрипт
+ python_executable = sys.executable
+
+ # Запускаем скрипт evaluate_chunking.py с нужными параметрами
+ cmd = [
+ python_executable, "scripts/evaluate_chunking.py",
+ "--data-folder", args.data_folder,
+ "--model-name", model_name,
+ "--dataset-path", args.dataset_path,
+ "--output-dir", args.output_dir,
+ "--words-per-chunk", str(words),
+ "--overlap-words", str(overlap),
+ "--similarity-threshold", str(args.similarity_threshold),
+ "--device", args.device,
+ "--force-recompute" # Принудительно пересчитываем эмбеддинги
+ ]
+
+ # Специальная обработка для модели ai-forever/FRIDA
+ if model_name == "ai-forever/FRIDA":
+ cmd.append("--use-sentence-transformers") # Добавляем флаг для использования sentence_transformers
+
+ print(f"\n{'='*80}")
+ print(f"Запуск эксперимента:")
+ print(f" Интерпретатор Python: {python_executable}")
+ print(f" Модель: {model_name}")
+ print(f" Чанкинг: {description} (words={words}, overlap={overlap})")
+ print(f" Порог для нечеткого сравнения: {args.similarity_threshold}")
+ print(f" Устройство: {args.device}")
+ print(f" Результаты будут сохранены в: {results_path}")
+ print(f" Лог: {log_path}")
+ print(f"{'='*80}\n")
+
+ # Запись информации в лог
+ with open(log_path, "w", encoding="utf-8") as log_file:
+ log_file.write(f"Эксперимент запущен в: {datetime.now()}\n")
+ log_file.write(f"Интерпретатор Python: {python_executable}\n")
+ log_file.write(f"Команда: {' '.join(cmd)}\n\n")
+
+ start_time = time.time()
+
+ # Запускаем процесс и перенаправляем вывод в файл лога
+ process = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ text=True,
+ bufsize=1 # Построчная буферизация
+ )
+
+ # Читаем вывод процесса
+ for line in process.stdout:
+ print(line, end="") # Выводим в консоль
+ log_file.write(line) # Записываем в файл лога
+
+ # Ждем завершения процесса
+ process.wait()
+
+ end_time = time.time()
+ duration = end_time - start_time
+
+ # Записываем информацию о завершении
+ log_file.write(f"\nЭксперимент завершен в: {datetime.now()}\n")
+ log_file.write(f"Длительность: {duration:.2f} секунд ({duration/60:.2f} минут)\n")
+ log_file.write(f"Код возврата: {process.returncode}\n")
+
+ if process.returncode == 0:
+ print(f"Эксперимент успешно завершен за {duration/60:.2f} минут")
+ else:
+ print(f"Эксперимент завершился с ошибкой (код {process.returncode})")
+
+
+def main():
+ """Основная функция скрипта."""
+ args = parse_args()
+
+ # Создаем output_dir, если он не существует
+ os.makedirs(args.output_dir, exist_ok=True)
+
+ # Получаем список моделей для запуска
+ models_to_run = [args.model] if args.model else MODELS
+
+ # Получаем список конфигураций чанкинга для запуска
+ chunking_configs = [CHUNKING_PARAMS[args.chunking_index]] if args.chunking_index is not None else CHUNKING_PARAMS
+
+ start_time_all = time.time()
+ total_experiments = len(models_to_run) * len(chunking_configs)
+ completed_experiments = 0
+
+ print(f"Запуск {total_experiments} экспериментов...")
+
+ # Изменен порядок: сначала идём по стратегиям, затем по моделям
+ for chunking_config in chunking_configs:
+ print(f"\n=== Стратегия чанкинга: {chunking_config['description']} (words={chunking_config['words']}, overlap={chunking_config['overlap']}) ===\n")
+
+ for model in models_to_run:
+ # Запускаем эксперимент
+ run_experiment(model, chunking_config, args)
+
+ completed_experiments += 1
+ remaining_experiments = total_experiments - completed_experiments
+
+ if remaining_experiments > 0:
+ print(f"Завершено {completed_experiments}/{total_experiments} экспериментов. Осталось: {remaining_experiments}")
+
+ end_time_all = time.time()
+ total_duration = end_time_all - start_time_all
+
+ print(f"\nВсе эксперименты завершены за {total_duration/60:.2f} минут")
+ print(f"Результаты сохранены в {args.output_dir}")
+ print(f"Логи сохранены в {args.log_dir}")
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/lib/extractor/scripts/search_api.py b/lib/extractor/scripts/search_api.py
new file mode 100644
index 0000000000000000000000000000000000000000..e621ad444594822effb3fdc578872e1c0ec31c0f
--- /dev/null
+++ b/lib/extractor/scripts/search_api.py
@@ -0,0 +1,748 @@
+#!/usr/bin/env python
+"""
+Скрипт для поиска по векторизованным документам через API.
+
+Этот скрипт:
+1. Считывает все документы из заданной папки с помощью UniversalParser
+2. Чанкит каждый документ через Destructurer с fixed_size-стратегией
+3. Векторизует поле in_search_text через BGE-модель
+4. Поднимает FastAPI с двумя эндпоинтами:
+ - /search/entities - возвращает найденные сущности списком словарей
+ - /search/text - возвращает полноценный собранный текст
+"""
+
+import logging
+import os
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import uvicorn
+from fastapi import FastAPI, Query
+from ntr_fileparser import UniversalParser
+from pydantic import BaseModel
+from sklearn.metrics.pairwise import cosine_similarity
+from transformers import AutoModel, AutoTokenizer
+
+from ntr_text_fragmentation.chunking.specific_strategies.fixed_size_chunking import \
+ FixedSizeChunkingStrategy
+from ntr_text_fragmentation.core.destructurer import Destructurer
+from ntr_text_fragmentation.core.entity_repository import \
+ InMemoryEntityRepository
+from ntr_text_fragmentation.core.injection_builder import InjectionBuilder
+from ntr_text_fragmentation.models.linker_entity import LinkerEntity
+
+# Константы
+DOCS_FOLDER = "../data/docs" # Путь к папке с документами
+MODEL_NAME = "BAAI/bge-m3" # Название модели для векторизации
+BATCH_SIZE = 16 # Размер батча для векторизации
+DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" # Устройство для вычислений
+MAX_ENTITIES = 100 # Максимальное количество возвращаемых сущностей
+WORDS_PER_CHUNK = 50 # Количество слов в чанке для fixed_size стратегии
+OVERLAP_WORDS = 25 # Количество слов перекрытия для fixed_size стратегии
+
+# Пути к кэшированным файлам
+CACHE_DIR = "../data/cache" # Путь к папке с кэшированными данными
+ENTITIES_CSV = os.path.join(CACHE_DIR, "entities.csv") # Путь к CSV с сущностями
+EMBEDDINGS_NPY = os.path.join(CACHE_DIR, "embeddings.npy") # Путь к массиву эмбеддингов
+
+# Инициализация FastAPI
+app = FastAPI(title="Документный поиск API",
+ description="API для поиска по векторизованным документам")
+
+# Глобальные переменные для хранения данных
+entities_df = None
+entity_embeddings = None
+model = None
+tokenizer = None
+entity_repository = None
+injection_builder = None
+
+
+class EntityResponse(BaseModel):
+ """Модель ответа для сущностей."""
+ id: str
+ name: str
+ text: str
+ type: str
+ score: float
+ doc_name: Optional[str] = None
+ metadata: Optional[Dict] = None
+
+
+class TextResponse(BaseModel):
+ """Модель ответа для собранного текста."""
+ text: str
+ entities_count: int
+
+
+class TextsResponse(BaseModel):
+ """Модель ответа для списка текстов."""
+ texts: List[str]
+ entities_count: int
+
+
+def setup_logging() -> None:
+ """Настройка логгирования."""
+ logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+ )
+
+
+def load_documents(folder_path: str) -> Dict:
+ """
+ Загружает все документы из указанной папки.
+
+ Args:
+ folder_path: Путь к папке с документами
+
+ Returns:
+ Словарь {имя_файла: parsed_document}
+ """
+ logging.info(f"Чтение документов из {folder_path}...")
+ parser = UniversalParser()
+ documents = {}
+
+ # Проверка существования папки
+ if not os.path.exists(folder_path):
+ logging.error(f"Папка {folder_path} не существует!")
+ return {}
+
+ for file_path in Path(folder_path).glob("**/*.docx"):
+ try:
+ doc_name = file_path.stem
+ logging.info(f"Обработка документа: {doc_name}")
+ documents[doc_name] = parser.parse_by_path(str(file_path))
+ except Exception as e:
+ logging.error(f"Ошибка при чтении файла {file_path}: {e}")
+
+ logging.info(f"Загружено {len(documents)} документов.")
+ return documents
+
+
+def process_documents(documents: Dict) -> List[LinkerEntity]:
+ """
+ Обрабатывает документы, применяя fixed_size стратегию чанкинга.
+
+ Args:
+ documents: Словарь с распарсенными документами
+
+ Returns:
+ Список сущностей из всех документов
+ """
+ logging.info("Применение fixed_size стратегии чанкинга ко всем документам...")
+
+ all_entities = []
+
+ for doc_name, document in documents.items():
+ try:
+ # Создаем Destructurer с fixed_size стратегией
+ destructurer = Destructurer(
+ document,
+ strategy_name="fixed_size",
+ words_per_chunk=WORDS_PER_CHUNK,
+ overlap_words=OVERLAP_WORDS
+ )
+
+ # Получаем сущности
+ doc_entities = destructurer.destructure()
+
+ # Добавляем имя документа в метаданные всех сущностей
+ for entity in doc_entities:
+ if not hasattr(entity, 'metadata') or entity.metadata is None:
+ entity.metadata = {}
+ entity.metadata['doc_name'] = doc_name
+
+ all_entities.extend(doc_entities)
+ logging.info(f"Документ {doc_name}: получено {len(doc_entities)} сущностей")
+
+ except Exception as e:
+ logging.error(f"Ошибка при обработке документа {doc_name}: {e}")
+
+ logging.info(f"Всего получено {len(all_entities)} сущностей из всех документов")
+ return all_entities
+
+
+def entities_to_dataframe(entities: List[LinkerEntity]) -> pd.DataFrame:
+ """
+ Преобразует список сущностей в DataFrame для удобной работы.
+
+ Args:
+ entities: Список сущностей
+
+ Returns:
+ DataFrame с данными сущностей
+ """
+ data = []
+
+ for entity in entities:
+ # Получаем имя документа из метаданных
+ doc_name = entity.metadata.get('doc_name', '') if hasattr(entity, 'metadata') and entity.metadata else ''
+
+ # Базовые поля для всех типов сущностей
+ entity_dict = {
+ "id": str(entity.id),
+ "type": entity.type,
+ "name": entity.name,
+ "text": entity.text,
+ "in_search_text": entity.in_search_text,
+ "doc_name": doc_name,
+ "source_id": entity.source_id if hasattr(entity, 'source_id') else None,
+ "target_id": entity.target_id if hasattr(entity, 'target_id') else None,
+ "metadata": entity.metadata if hasattr(entity, 'metadata') else {},
+ }
+
+ data.append(entity_dict)
+
+ df = pd.DataFrame(data)
+ return df
+
+
+def setup_model_and_tokenizer():
+ """
+ Инициализирует модель и токенизатор для векторизации.
+
+ Returns:
+ Кортеж (модель, токенизатор)
+ """
+ global model, tokenizer
+
+ logging.info(f"Загрузка модели {MODEL_NAME} на устройство {DEVICE}...")
+
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+ model = AutoModel.from_pretrained(MODEL_NAME).to(DEVICE)
+ model.eval()
+
+ return model, tokenizer
+
+
+def _average_pool(
+ last_hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor
+) -> torch.Tensor:
+ """
+ Расчёт усредненного эмбеддинга по всем токенам
+
+ Args:
+ last_hidden_states: Матрица эмбеддингов отдельных токенов
+ attention_mask: Маска, чтобы не учитывать при усреднении пустые токены
+
+ Returns:
+ Усредненный эмбеддинг
+ """
+ last_hidden = last_hidden_states.masked_fill(
+ ~attention_mask[..., None].bool(), 0.0
+ )
+ return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+
+def get_embeddings(texts: List[str]) -> np.ndarray:
+ """
+ Получает эмбеддинги для списка текстов.
+
+ Args:
+ texts: Список текстов для векторизации
+
+ Returns:
+ Массив эмбеддингов
+ """
+ global model, tokenizer
+
+ # Проверяем, что модель и токенизатор инициализированы
+ if model is None or tokenizer is None:
+ model, tokenizer = setup_model_and_tokenizer()
+
+ all_embeddings = []
+
+ for i in range(0, len(texts), BATCH_SIZE):
+ batch_texts = texts[i:i+BATCH_SIZE]
+
+ # Фильтруем None и пустые строки
+ batch_texts = [text for text in batch_texts if text]
+
+ if not batch_texts:
+ continue
+
+ # Токенизация с обрезкой и padding
+ encoding = tokenizer(
+ batch_texts,
+ padding=True,
+ truncation=True,
+ max_length=512,
+ return_tensors="pt"
+ ).to(DEVICE)
+
+ # Получаем эмбеддинги с average pooling
+ with torch.no_grad():
+ outputs = model(**encoding)
+ embeddings = _average_pool(outputs.last_hidden_state, encoding["attention_mask"])
+ all_embeddings.append(embeddings.cpu().numpy())
+
+ if not all_embeddings:
+ return np.array([])
+
+ return np.vstack(all_embeddings)
+
+
+def init_entity_repository_and_builder(entities: List[LinkerEntity]):
+ """
+ Инициализирует хранилище сущностей и сборщик инъекций.
+
+ Args:
+ entities: Список сущностей
+ """
+ global entity_repository, injection_builder
+
+ # Создаем хранилище сущностей
+ entity_repository = InMemoryEntityRepository(entities)
+
+ # Добавляем метод get_entity_by_id в InMemoryEntityRepository
+ # Это временное решение, в идеале нужно добавить этот метод в сам класс
+ def get_entity_by_id(self, entity_id):
+ """Получает сущность по ID"""
+ for entity in self.entities:
+ if str(entity.id) == entity_id:
+ return entity
+ return None
+
+ # Добавляем метод в класс
+ InMemoryEntityRepository.get_entity_by_id = get_entity_by_id
+
+ # Создаем сборщик инъекций
+ injection_builder = InjectionBuilder(repository=entity_repository)
+
+ # Регистрируем стратегию
+ injection_builder.register_strategy("fixed_size", FixedSizeChunkingStrategy)
+
+
+def search_entities(query: str, top_n: int = MAX_ENTITIES) -> List[Dict]:
+ """
+ Ищет сущности по запросу на основе косинусной близости.
+
+ Args:
+ query: Поисковый запрос
+ top_n: Максимальное количество возвращаемых сущностей
+
+ Returns:
+ Список найденных сущностей с их скорами
+ """
+ global entities_df, entity_embeddings
+
+ # Проверяем наличие данных
+ if entities_df is None or entity_embeddings is None:
+ logging.error("Данные не инициализированы. Запустите сначала prepare_data().")
+ return []
+
+ # Векторизуем запрос
+ query_embedding = get_embeddings([query])
+
+ if query_embedding.size == 0:
+ return []
+
+ # Считаем косинусную близость
+ similarities = cosine_similarity(query_embedding, entity_embeddings)[0]
+
+ # Получаем индексы топ-N сущностей
+ top_indices = np.argsort(similarities)[-top_n:][::-1]
+
+ # Фильтруем сущности, которые используются для поиска
+ search_df = entities_df.copy()
+ search_df = search_df[search_df['in_search_text'].notna()]
+
+ # Если после фильтрации нет данных, возвращаем пустой список
+ if search_df.empty:
+ return []
+
+ # Получаем топ-N сущностей
+ results = []
+
+ for idx in top_indices:
+ if idx >= len(search_df):
+ continue
+
+ entity = search_df.iloc[idx]
+ similarity = similarities[idx]
+
+ # Создаем результат
+ result = {
+ "id": entity["id"],
+ "name": entity["name"],
+ "text": entity["text"],
+ "type": entity["type"],
+ "score": float(similarity),
+ "doc_name": entity["doc_name"],
+ "metadata": entity["metadata"]
+ }
+
+ results.append(result)
+
+ return results
+
+
+@app.get("/search/entities", response_model=List[EntityResponse])
+async def api_search_entities(
+ query: str = Query(..., description="Поисковый запрос"),
+ limit: int = Query(MAX_ENTITIES, description="Максимальное количество результатов")
+):
+ """
+ Эндпоинт для поиска сущностей по запросу.
+
+ Args:
+ query: Поисковый запрос
+ limit: Максимальное количество результатов
+
+ Returns:
+ Список найденных сущностей
+ """
+ results = search_entities(query, limit)
+ return results
+
+
+@app.get("/search/text", response_model=TextResponse)
+async def api_search_text(
+ query: str = Query(..., description="Поисковый запрос"),
+ limit: int = Query(MAX_ENTITIES, description="Максимальное количество учитываемых сущностей")
+):
+ """
+ Эндпоинт для поиска и сборки полного текста по запросу.
+
+ Args:
+ query: Поисковый запрос
+ limit: Максимальное количество учитываемых сущностей
+
+ Returns:
+ Собранный текст и количество использованных сущностей
+ """
+ global injection_builder
+
+ # Проверяем наличие сборщика инъекций
+ if injection_builder is None:
+ logging.error("Сборщик инъекций не инициализирован.")
+ return {"text": "", "entities_count": 0}
+
+ # Получаем найденные сущности
+ entity_results = search_entities(query, limit)
+
+ if not entity_results:
+ return {"text": "", "entities_count": 0}
+
+ # Получаем список ID сущностей
+ entity_ids = [str(result["id"]) for result in entity_results]
+
+ # Собираем текст, используя напрямую ID
+ try:
+ assembled_text = injection_builder.build(entity_ids)
+ print('Всё ок прошло вроде бы')
+ return {"text": assembled_text, "entities_count": len(entity_ids)}
+ except ImportError as e:
+ # Обработка ошибки импорта модулей для работы с изображениями
+ logging.error(f"Ошибка импорта при сборке текста: {e}")
+ # Альтернативная сборка текста без использования injection_builder
+ simple_text = "\n\n".join([result["text"] for result in entity_results if result.get("text")])
+ return {"text": simple_text, "entities_count": len(entity_ids)}
+ except Exception as e:
+ logging.error(f"Ошибка при сборке текста: {e}")
+ return {"text": "", "entities_count": 0}
+
+
+@app.get("/search/texts", response_model=TextsResponse)
+async def api_search_texts(
+ query: str = Query(..., description="Поисковый запрос"),
+ limit: int = Query(MAX_ENTITIES, description="Максимальное количество результатов")
+):
+ """
+ Эндпоинт для поиска списка текстов сущностей по запросу.
+
+ Args:
+ query: Поисковый запрос
+ limit: Максимальное количество результатов
+
+ Returns:
+ Список текстов найденных сущностей и их количество
+ """
+ # Получаем найденные сущности
+ entity_results = search_entities(query, limit)
+
+ if not entity_results:
+ return {"texts": [], "entities_count": 0}
+
+ # Извлекаем тексты из результатов
+ texts = [result["text"] for result in entity_results if result.get("text")]
+
+ return {"texts": texts, "entities_count": len(texts)}
+
+
+@app.get("/search/text_test", response_model=TextResponse)
+async def api_search_text_test(
+ query: str = Query(..., description="Поисковый запрос"),
+ limit: int = Query(MAX_ENTITIES, description="Максимальное количество учитываемых сущностей")
+):
+ """
+ Тестовый эндпоинт для поиска и сборки текста с использованием подхода из test_chunking_visualization.py.
+
+ Args:
+ query: Поисковый запрос
+ limit: Максимальное количество учитываемых сущностей
+
+ Returns:
+ Собранный текст и количество использованных сущностей
+ """
+ global entity_repository, injection_builder
+
+ # Проверяем наличие репозитория и сборщика инъекций
+ if entity_repository is None or injection_builder is None:
+ logging.error("Репозиторий или сборщик инъекций не инициализированы.")
+ return {"text": "", "entities_count": 0}
+
+ # Получаем найденные сущности
+ entity_results = search_entities(query, limit)
+
+ if not entity_results:
+ return {"text": "", "entities_count": 0}
+
+ try:
+ # Получаем объекты сущностей из репозитория по ID
+ entity_ids = [result["id"] for result in entity_results]
+ entities = []
+
+ for entity_id in entity_ids:
+ entity = entity_repository.get_entity_by_id(entity_id)
+ if entity:
+ entities.append(entity)
+
+ logging.info(f"Найдено {len(entities)} объектов сущностей по ID")
+
+ if not entities:
+ logging.error("Не удалось найти сущности в репозитории")
+ # Собираем простой текст из результатов поиска
+ simple_text = "\n\n".join([result["text"] for result in entity_results if result.get("text")])
+ return {"text": simple_text, "entities_count": len(entity_results)}
+
+ # Собираем текст, как в test_chunking_visualization.py
+ assembled_text = injection_builder.build(entities) # Передаем сами объекты
+
+ return {"text": assembled_text, "entities_count": len(entities)}
+ except Exception as e:
+ logging.error(f"Ошибка при сборке текста: {e}", exc_info=True)
+ # Запасной вариант - просто соединяем тексты
+ fallback_text = "\n\n".join([result["text"] for result in entity_results if result.get("text")])
+ return {"text": fallback_text, "entities_count": len(entity_results)}
+
+
+def save_entities_to_csv(entities: List[LinkerEntity], csv_path: str) -> None:
+ """
+ Сохраняет сущности в CSV файл.
+
+ Args:
+ entities: Список сущностей
+ csv_path: Путь для сохранения CSV файла
+ """
+ logging.info(f"Сохранение {len(entities)} сущностей в {csv_path}")
+
+ # Создаем директорию, если она не существует
+ os.makedirs(os.path.dirname(csv_path), exist_ok=True)
+
+ # Преобразуем сущности в DataFrame и сохраняем
+ df = entities_to_dataframe(entities)
+ df.to_csv(csv_path, index=False)
+
+ logging.info(f"Сохранено {len(entities)} сущностей в {csv_path}")
+
+
+def load_entities_from_csv(csv_path: str) -> List[LinkerEntity]:
+ """
+ Загружает сущности из CSV файла.
+
+ Args:
+ csv_path: Путь к CSV файлу
+
+ Returns:
+ Список сущностей
+ """
+ logging.info(f"Загрузка сущностей из {csv_path}")
+
+ if not os.path.exists(csv_path):
+ logging.error(f"Файл {csv_path} не найден")
+ return []
+
+ df = pd.read_csv(csv_path)
+ entities = []
+
+ for _, row in df.iterrows():
+ # Обработка метаданных
+ metadata = row.get("metadata", {})
+ if isinstance(metadata, str):
+ try:
+ metadata = eval(metadata) if metadata and not pd.isna(metadata) else {}
+ except:
+ metadata = {}
+
+ # Общие поля для всех типов сущностей
+ common_args = {
+ "id": row["id"],
+ "name": row["name"] if not pd.isna(row.get("name", "")) else "",
+ "text": row["text"] if not pd.isna(row.get("text", "")) else "",
+ "metadata": metadata,
+ "type": row["type"],
+ }
+
+ # Добавляем in_search_text, если он есть
+ if "in_search_text" in row and not pd.isna(row["in_search_text"]):
+ common_args["in_search_text"] = row["in_search_text"]
+
+ # Добавляем поля связи, если они есть
+ if "source_id" in row and not pd.isna(row["source_id"]):
+ common_args["source_id"] = row["source_id"]
+ common_args["target_id"] = row["target_id"]
+ if "number_in_relation" in row and not pd.isna(row["number_in_relation"]):
+ common_args["number_in_relation"] = int(row["number_in_relation"])
+
+ entity = LinkerEntity(**common_args)
+ entities.append(entity)
+
+ logging.info(f"Загружено {len(entities)} сущностей из {csv_path}")
+ return entities
+
+
+def save_embeddings(embeddings: np.ndarray, file_path: str) -> None:
+ """
+ Сохраняет эмбеддинги в numpy файл.
+
+ Args:
+ embeddings: Массив эмбеддингов
+ file_path: Путь для сохранения файла
+ """
+ logging.info(f"Сохранение эмбеддингов размером {embeddings.shape} в {file_path}")
+
+ # Создаем директорию, если она не существует
+ os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+ # Сохраняем эмбеддинги
+ np.save(file_path, embeddings)
+
+ logging.info(f"Эмбеддинги сохранены в {file_path}")
+
+
+def load_embeddings(file_path: str) -> np.ndarray:
+ """
+ Загружает эмбеддинги из numpy файла.
+
+ Args:
+ file_path: Путь к файлу
+
+ Returns:
+ Массив эмбеддингов
+ """
+ logging.info(f"Загрузка эмбеддингов из {file_path}")
+
+ if not os.path.exists(file_path):
+ logging.error(f"Файл {file_path} не найден")
+ return np.array([])
+
+ embeddings = np.load(file_path)
+
+ logging.info(f"Загружены эмбеддинги размером {embeddings.shape}")
+ return embeddings
+
+
+def prepare_data():
+ """
+ Подготавливает все необходимые данные для API.
+ """
+ global entities_df, entity_embeddings, entity_repository, injection_builder
+
+ # Проверяем наличие кэшированных данных
+ cache_exists = os.path.exists(ENTITIES_CSV) and os.path.exists(EMBEDDINGS_NPY)
+
+ if cache_exists:
+ logging.info("Найдены кэшированные данные, загружаем их")
+
+ # Загружаем сущности из CSV
+ entities = load_entities_from_csv(ENTITIES_CSV)
+
+ if not entities:
+ logging.error("Не удалось загрузить сущности из кэша, генерируем заново")
+ cache_exists = False
+ else:
+ # Преобразуем сущности в DataFrame
+ entities_df = entities_to_dataframe(entities)
+
+ # Загружаем эмбеддинги
+ entity_embeddings = load_embeddings(EMBEDDINGS_NPY)
+
+ if entity_embeddings.size == 0:
+ logging.error("Не удалось загрузить эмбеддинги из кэша, генерируем заново")
+ cache_exists = False
+ else:
+ # Инициализируем хранилище и сборщик
+ init_entity_repository_and_builder(entities)
+ logging.info("Данные успешно загружены из кэша")
+
+ # Если кэшированных данных нет или их не удалось загрузить, генерируем заново
+ if not cache_exists:
+ logging.info("Кэшированные данные не найдены или не могут быть загружены, обрабатываем документы")
+
+ # Загружаем и обрабатываем документы
+ documents = load_documents(DOCS_FOLDER)
+
+ if not documents:
+ logging.error(f"Не найдено документов в папке {DOCS_FOLDER}")
+ return
+
+ # Получаем сущности из всех документов
+ all_entities = process_documents(documents)
+
+ if not all_entities:
+ logging.error("Не получено сущностей из документов")
+ return
+
+ # Преобразуем сущности в DataFrame
+ entities_df = entities_to_dataframe(all_entities)
+
+ # Инициализируем хранилище и сборщик
+ init_entity_repository_and_builder(all_entities)
+
+ # Фильтруем только сущности для поиска
+ search_df = entities_df[entities_df['in_search_text'].notna()]
+
+ if search_df.empty:
+ logging.error("Нет сущностей для поиска с in_search_text")
+ return
+
+ # Векторизуем тексты сущностей
+ search_texts = search_df['in_search_text'].tolist()
+ entity_embeddings = get_embeddings(search_texts)
+
+ logging.info(f"Подготовлено {len(search_df)} сущностей для поиска")
+ logging.info(f"Размер эмбеддингов: {entity_embeddings.shape}")
+
+ # Сохраняем данные в кэш для последующего использования
+ save_entities_to_csv(all_entities, ENTITIES_CSV)
+ save_embeddings(entity_embeddings, EMBEDDINGS_NPY)
+ logging.info("Данные сохранены в кэш для последующего использования")
+
+ # Вывод итоговой информации (независимо от источника данных)
+ logging.info(f"Подготовка данных завершена. Готово к использованию {entity_embeddings.shape[0]} сущностей")
+
+
+@app.on_event("startup")
+async def startup_event():
+ """Запускается при старте приложения."""
+ setup_logging()
+ prepare_data()
+
+
+def main():
+ """Основная функция для запуска скрипта вручную."""
+ setup_logging()
+ prepare_data()
+
+ # Запуск Uvicorn сервера
+ uvicorn.run(app, host="0.0.0.0", port=8017)
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/lib/extractor/scripts/test_chunking_visualization.py b/lib/extractor/scripts/test_chunking_visualization.py
new file mode 100644
index 0000000000000000000000000000000000000000..2e0a639ebc2e4906827d6d2eb66cfe34b3e6583e
--- /dev/null
+++ b/lib/extractor/scripts/test_chunking_visualization.py
@@ -0,0 +1,235 @@
+#!/usr/bin/env python
+"""
+Скрипт для визуального тестирования процесса чанкинга и сборки документа.
+
+Этот скрипт:
+1. Считывает test_input/test.docx с помощью UniversalParser
+2. Чанкит документ через Destructurer с fixed_size-стратегией
+3. Сохраняет результат чанкинга в test_output/test.csv
+4. Выбирает 20-30 случайных чанков из CSV
+5. Создает InjectionBuilder с InMemoryEntityRepository
+6. Собирает текст из выбранных чанков
+7. Сохраняет результат в test_output/test_builded.txt
+"""
+
+import logging
+import os
+import random
+from pathlib import Path
+from typing import List
+
+import pandas as pd
+from ntr_fileparser import UniversalParser
+
+from ntr_text_fragmentation.chunking.specific_strategies.fixed_size_chunking import \
+ FixedSizeChunkingStrategy
+from ntr_text_fragmentation.core.destructurer import Destructurer
+from ntr_text_fragmentation.core.entity_repository import \
+ InMemoryEntityRepository
+from ntr_text_fragmentation.core.injection_builder import InjectionBuilder
+from ntr_text_fragmentation.models.linker_entity import LinkerEntity
+
+
+def setup_logging() -> None:
+ """Настройка логгирования."""
+ logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+ )
+
+
+def ensure_directories() -> None:
+ """Проверка наличия необходимых директорий."""
+ for directory in ["test_input", "test_output"]:
+ Path(directory).mkdir(parents=True, exist_ok=True)
+
+
+def save_entities_to_csv(entities: List[LinkerEntity], csv_path: str) -> None:
+ """
+ Сохраняет сущности в CSV файл.
+
+ Args:
+ entities: Список сущностей
+ csv_path: Путь для сохранения CSV файла
+ """
+ data = []
+ for entity in entities:
+ # Базовые поля для всех типов сущностей
+ entity_dict = {
+ "id": str(entity.id),
+ "type": entity.type,
+ "name": entity.name,
+ "text": entity.text,
+ "metadata": str(entity.metadata),
+ "in_search_text": entity.in_search_text,
+ "source_id": entity.source_id,
+ "target_id": entity.target_id,
+ "number_in_relation": entity.number_in_relation,
+ }
+
+ data.append(entity_dict)
+
+ df = pd.DataFrame(data)
+ df.to_csv(csv_path, index=False)
+ logging.info(f"Сохранено {len(entities)} сущностей в {csv_path}")
+
+
+def load_entities_from_csv(csv_path: str) -> List[LinkerEntity]:
+ """
+ Загружает сущности из CSV файла.
+
+ Args:
+ csv_path: Путь к CSV файлу
+
+ Returns:
+ Список сущностей
+ """
+ df = pd.read_csv(csv_path)
+ entities = []
+
+ for _, row in df.iterrows():
+ # Обработка метаданных
+ metadata_str = row.get("metadata", "{}")
+ try:
+ metadata = (
+ eval(metadata_str) if metadata_str and not pd.isna(metadata_str) else {}
+ )
+ except:
+ metadata = {}
+
+ # Общие поля для всех типов сущностей
+ common_args = {
+ "id": row["id"],
+ "name": row["name"] if not pd.isna(row.get("name", "")) else "",
+ "text": row["text"] if not pd.isna(row.get("text", "")) else "",
+ "metadata": metadata,
+ "in_search_text": row["in_search_text"],
+ "type": row["type"],
+ }
+
+ # Добавляем поля связи, если они есть
+ if not pd.isna(row.get("source_id", "")):
+ common_args["source_id"] = row["source_id"]
+ common_args["target_id"] = row["target_id"]
+ if not pd.isna(row.get("number_in_relation", "")):
+ common_args["number_in_relation"] = int(row["number_in_relation"])
+
+ entity = LinkerEntity(**common_args)
+ entities.append(entity)
+
+ logging.info(f"Загружено {len(entities)} сущностей из {csv_path}")
+ return entities
+
+
+def main() -> None:
+ """Основная функция скрипта."""
+ setup_logging()
+ ensure_directories()
+
+ # Пути к файлам
+ input_doc_path = "test_input/test.docx"
+ output_csv_path = "test_output/test.csv"
+ output_text_path = "test_output/test_builded.txt"
+
+ # Проверка наличия входного файла
+ if not os.path.exists(input_doc_path):
+ logging.error(f"Файл {input_doc_path} не найден!")
+ return
+
+ logging.info(f"Парсинг документа {input_doc_path}")
+
+ try:
+ # Шаг 1: Парсинг документа дважды, как если бы это были два разных документа
+ parser = UniversalParser()
+ document1 = parser.parse_by_path(input_doc_path)
+ document2 = parser.parse_by_path(input_doc_path)
+
+ # Меняем название второго документа, чтобы отличить его
+ document2.name = document2.name + "_copy" if document2.name else "copy_doc"
+
+ # Шаг 2: Чанкинг обоих документов с использованием fixed_size-стратегии
+ all_entities = []
+
+ # Обработка первого документа
+ destructurer1 = Destructurer(
+ document1, strategy_name="fixed_size", words_per_chunk=50, overlap_words=25
+ )
+ logging.info("Начало процесса чанкинга первого документа")
+ entities1 = destructurer1.destructure()
+
+ # Добавляем метаданные о документе к каждой сущности
+ for entity in entities1:
+ if not hasattr(entity, 'metadata') or entity.metadata is None:
+ entity.metadata = {}
+ entity.metadata['doc_name'] = "document1"
+
+ logging.info(f"Получено {len(entities1)} сущностей из первого документа")
+ all_entities.extend(entities1)
+
+ # Обработка второго документа
+ destructurer2 = Destructurer(
+ document2, strategy_name="fixed_size", words_per_chunk=50, overlap_words=25
+ )
+ logging.info("Начало процесса чанкинга второго документа")
+ entities2 = destructurer2.destructure()
+
+ # Добавляем метаданные о документе к каждой сущности
+ for entity in entities2:
+ if not hasattr(entity, 'metadata') or entity.metadata is None:
+ entity.metadata = {}
+ entity.metadata['doc_name'] = "document2"
+
+ logging.info(f"Получено {len(entities2)} сущностей из второго документа")
+ all_entities.extend(entities2)
+
+ logging.info(f"Всего получено {len(all_entities)} сущностей из обоих документов")
+
+ # Шаг 3: Сохранение результатов чанкинга в CSV
+ save_entities_to_csv(all_entities, output_csv_path)
+
+ # Шаг 4: Загрузка сущностей из CSV и выбор случайных чанков
+ loaded_entities = load_entities_from_csv(output_csv_path)
+
+ # Фильтрация только чанков
+ chunks = [e for e in loaded_entities if e.in_search_text is not None]
+
+ # Выбор случайных чанков (от 20 до 30)
+ num_chunks_to_select = min(random.randint(20, 30), len(chunks))
+ selected_chunks = random.sample(chunks, num_chunks_to_select)
+
+ logging.info(f"Выбрано {len(selected_chunks)} случайных чанков для сборки")
+
+ # Дополнительная статистика по документам
+ doc1_chunks = [c for c in selected_chunks if hasattr(c, 'metadata') and c.metadata.get('doc_name') == "document1"]
+ doc2_chunks = [c for c in selected_chunks if hasattr(c, 'metadata') and c.metadata.get('doc_name') == "document2"]
+ logging.info(f"Из них {len(doc1_chunks)} чанков из первого документа и {len(doc2_chunks)} из второго")
+
+ # Шаг 5: Создание InjectionBuilder с InMemoryEntityRepository
+ repository = InMemoryEntityRepository(loaded_entities)
+ builder = InjectionBuilder(repository=repository)
+
+ # Регистрация стратегии
+ builder.register_strategy("fixed_size", FixedSizeChunkingStrategy)
+
+ # Шаг 6: Сборка текста из выбранных чанков
+ logging.info("Начало сборки текста из выбранных чанков")
+ assembled_text = builder.build(selected_chunks)
+
+ # Шаг 7: Сохранение результата в файл
+ with open(output_text_path, "w", encoding="utf-8") as f:
+ f.write(assembled_text)
+
+ logging.info(f"Результат сборки сохранен в {output_text_path}")
+
+ # Вывод статистики
+ logging.info(f"Общее количество сущностей: {len(loaded_entities)}")
+ logging.info(f"Количество чанков: {len(chunks)}")
+ logging.info(f"Выбрано для сборки: {len(selected_chunks)}")
+ logging.info(f"Длина собранного текста: {len(assembled_text)} символов")
+
+ except Exception as e:
+ logging.error(f"Произошла ошибка: {e}", exc_info=True)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/lib/extractor/tests/__init__.py b/lib/extractor/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..523ea57ec87cd3769e12fb950b83d93563e6bddf
--- /dev/null
+++ b/lib/extractor/tests/__init__.py
@@ -0,0 +1,3 @@
+"""
+Пакет с тестами для ntr_text_fragmentation.
+"""
\ No newline at end of file
diff --git a/lib/extractor/tests/chunking/__init__.py b/lib/extractor/tests/chunking/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1361bb2a52eb29ed937f3df3633e79a5de0a9554
--- /dev/null
+++ b/lib/extractor/tests/chunking/__init__.py
@@ -0,0 +1,3 @@
+"""
+Тесты для компонентов чанкинга.
+"""
\ No newline at end of file
diff --git a/lib/extractor/tests/chunking/test_fixed_size_chunking.py b/lib/extractor/tests/chunking/test_fixed_size_chunking.py
new file mode 100644
index 0000000000000000000000000000000000000000..36b939e72d7d539f7b21ef6fc7f14eb02a6fe3db
--- /dev/null
+++ b/lib/extractor/tests/chunking/test_fixed_size_chunking.py
@@ -0,0 +1,334 @@
+from uuid import UUID
+
+import pytest
+from ntr_fileparser import ParsedDocument, ParsedTextBlock
+
+from ntr_text_fragmentation.chunking.specific_strategies.fixed_size_chunking import \
+ FixedSizeChunkingStrategy
+from ntr_text_fragmentation.models import DocumentAsEntity
+
+
+class TestFixedSizeChunkingStrategy:
+ """Набор тестов для проверки стратегии чанкинга фиксированного размера."""
+
+ @pytest.fixture
+ def sample_document(self):
+ """Фикстура для создания тестового документа."""
+ paragraphs = [
+ ParsedTextBlock(
+ text="Это первый параграф тестового документа. Он содержит два предложения."
+ ),
+ ParsedTextBlock(text="Это второй параграф с одним предложением."),
+ ParsedTextBlock(
+ text="Третий параграф. Содержит еще два предложения. И оно короткое."
+ ),
+ ]
+
+ return ParsedDocument(
+ name="test_document.txt", type="text", paragraphs=paragraphs
+ )
+
+ @pytest.fixture
+ def doc_entity(self):
+ """Фикстура для создания сущности документа."""
+ return DocumentAsEntity(
+ id=UUID('12345678-1234-5678-1234-567812345678'),
+ name="Тестовый документ",
+ text="",
+ metadata={"type": "text"},
+ type="Document",
+ )
+
+ @pytest.fixture
+ def large_document(self):
+ """Фикстура для создания большого тестового документа."""
+ paragraphs = [
+ ParsedTextBlock(
+ text="Это первый параграф большого документа. Он содержит несколько предложений разной длины."
+ ),
+ ParsedTextBlock(
+ text="Второй параграф начинается с короткого предложения. А затем идет длинное предложение, которое содержит много слов и должно быть разбито на несколько чанков, потому что оно не помещается в один чанк стандартного размера."
+ ),
+ ParsedTextBlock(
+ text="Третий параграф содержит несколько предложений. Каждое предложение имеет свою структуру. И все они должны корректно обрабатываться."
+ ),
+ ParsedTextBlock(
+ text="Четвертый параграф начинается с длинного предложения, которое также должно быть разбито на несколько чанков, так как оно содержит много слов и не помещается в один чанк стандартного размера. А затем идет короткое предложение."
+ ),
+ ParsedTextBlock(
+ text="Пятый параграф. Содержит разные предложения. С разной пунктуацией. И разной структурой."
+ ),
+ ParsedTextBlock(
+ text="Шестой параграф начинается с короткого предложения. Затем идет длинное предложение, которое должно быть разбито на несколько чанков, потому что оно содержит много слов и не помещается в один чанк стандартного размера. И заканчивается коротким предложением."
+ ),
+ ParsedTextBlock(
+ text="Седьмой параграф содержит несколько предложений разной длины. Каждое предложение имеет свою структуру. И все они должны корректно обрабатываться."
+ ),
+ ParsedTextBlock(
+ text="Восьмой параграф начинается с длинного предложения, которое также должно быть разбито на несколько чанков, так как оно содержит много слов и не помещается в один чанк стандартного размера. А затем идет короткое предложение."
+ ),
+ ParsedTextBlock(
+ text="Девятый параграф. Содержит разные предложения. С разной пунктуацией. И разной структурой."
+ ),
+ ParsedTextBlock(
+ text="Десятый параграф начинается с короткого предложения. Затем идет длинное предложение, которое должно быть разбито на несколько чанков, потому что оно содержит много слов и не помещается в один чанк стандартного размера. И заканчивается коротким предложением."
+ ),
+ ]
+
+ return ParsedDocument(
+ name="large_test_document.txt", type="text", paragraphs=paragraphs
+ )
+
+ def test_basic_chunking_and_dechunking(self, sample_document, doc_entity):
+ """Тест базового сценария нарезки и сборки документа."""
+ strategy = FixedSizeChunkingStrategy(words_per_chunk=10, overlap_words=2)
+
+ # Разбиваем документ
+ entities = strategy.chunk(sample_document, doc_entity)
+
+ # Выделяем только чанки
+ chunks = [e for e in entities if e.type == "FixedSizeChunk"]
+
+ # Собираем документ обратно
+ result_text = strategy.dechunk(chunks)
+
+ # Проверяем, что текст не пустой
+ assert result_text
+
+ # Проверяем, что все слова из оригинального документа присутствуют в результате
+ original_text = " ".join([p.text for p in sample_document.paragraphs])
+ original_words = set(original_text.split())
+ result_words = set(result_text.split())
+
+ # Все оригинальные слова должны быть в результате
+ assert original_words.issubset(result_words)
+
+ # Проверяем, что длина результата примерно равна длине исходного текста
+ assert abs(len(result_text.split()) - len(original_text.split())) < 5
+
+ def test_chunking_with_different_sentence_lengths(self, doc_entity):
+ """Тест нарезки документа с предложениями разной длины."""
+ # Создаем документ с предложениями разной длины
+ text = (
+ "Короткое предложение. "
+ "Это предложение средней длины с несколькими словами. "
+ "А это очень длинное предложение, которое содержит много слов и должно быть разбито на несколько чанков, "
+ "потому что оно не помещается в один чанк стандартного размера. "
+ "И снова короткое."
+ )
+ doc = ParsedDocument(
+ name="test_document.txt",
+ type="text",
+ paragraphs=[ParsedTextBlock(text=text)],
+ )
+
+ strategy = FixedSizeChunkingStrategy(words_per_chunk=15, overlap_words=5)
+
+ # Разбиваем документ
+ entities = strategy.chunk(doc, doc_entity)
+ chunks = [e for e in entities if e.type == "FixedSizeChunk"]
+
+ # Проверяем, что длинное предложение было разбито на несколько чанков
+ assert len(chunks) > 1
+
+ # Собираем документ обратно
+ result_text = strategy.dechunk(chunks)
+
+ # Проверяем корректность сборки
+ original_words = set(text.split())
+ result_words = set(result_text.split())
+ assert original_words.issubset(result_words)
+
+ # Проверяем, что все предложения сохранились
+ original_sentences = set(s.strip() for s in text.split('.'))
+ result_sentences = set(s.strip() for s in result_text.split('.'))
+ assert original_sentences.issubset(result_sentences)
+
+ def test_empty_document(self, doc_entity):
+ """Тест обработки пустого документа."""
+ doc = ParsedDocument(name="empty.txt", type="text", paragraphs=[])
+
+ strategy = FixedSizeChunkingStrategy()
+
+ # Разбиваем документ
+ entities = strategy.chunk(doc, doc_entity)
+ chunks = [e for e in entities if e.type == "FixedSizeChunk"]
+
+ # Проверяем, что чанков нет
+ assert len(chunks) == 0
+
+ # Проверяем, что сборка пустого документа возвращает пустую строку
+ result_text = strategy.dechunk(chunks)
+ assert result_text == ""
+
+ def test_special_characters_and_punctuation(self, doc_entity):
+ """Тест обработки текста со специальными символами и пунктуацией."""
+ text = (
+ "Текст с разными символами: !@#$%^&*(). "
+ "Скобки (внутри) и [квадратные]. "
+ "Кавычки «елочки» и \"прямые\". "
+ "Тире — и дефис-. "
+ "Многоточие... и запятые, в разных местах."
+ )
+ doc = ParsedDocument(
+ name="test_document.txt",
+ type="text",
+ paragraphs=[ParsedTextBlock(text=text)],
+ )
+
+ strategy = FixedSizeChunkingStrategy(words_per_chunk=10, overlap_words=2)
+
+ # Разбиваем документ
+ entities = strategy.chunk(doc, doc_entity)
+ chunks = [e for e in entities if e.type == "FixedSizeChunk"]
+
+ # Собираем документ обратно
+ result_text = strategy.dechunk(chunks)
+
+ # Проверяем, что все специальные символы сохранились
+ special_chars = set('!@#$%^&*()[]«»"—...')
+ result_chars = set(result_text)
+ assert special_chars.issubset(result_chars)
+
+ # Проверяем, что текст совпадает с оригиналом
+ assert result_text == text
+
+ def test_large_document_chunking(self, large_document, doc_entity):
+ """Тест нарезки и сборки большого документа с множеством параграфов."""
+ strategy = FixedSizeChunkingStrategy(words_per_chunk=20, overlap_words=5)
+
+ # Разбиваем документ
+ entities = strategy.chunk(large_document, doc_entity)
+ chunks = [e for e in entities if e.type == "FixedSizeChunk"]
+
+ # Проверяем, что документ был разбит на несколько чанков
+ assert len(chunks) > 1
+
+ # Собираем документ обратно
+ result_text = strategy.dechunk(chunks)
+
+ # Получаем оригинальный текст
+ original_paragraphs = [p.text for p in large_document.paragraphs]
+
+ # Проверяем, что все параграфы сохранились
+ result_paragraphs = result_text.split('\n')
+ assert len(result_paragraphs) == len(original_paragraphs)
+
+ # Проверяем, что каждый параграф совпадает с оригиналом
+ for orig, res in zip(original_paragraphs, result_paragraphs):
+ assert orig.strip() == res.strip()
+
+ def test_exact_text_comparison(self, sample_document, doc_entity):
+ """Тест точного сравнения текстов после нарезки и сборки."""
+ strategy = FixedSizeChunkingStrategy(words_per_chunk=10, overlap_words=2)
+
+ # Разбиваем документ
+ entities = strategy.chunk(sample_document, doc_entity)
+ chunks = [e for e in entities if e.type == "FixedSizeChunk"]
+
+ # Собираем документ обратно
+ result_text = strategy.dechunk(chunks)
+
+ # Получаем оригинальный текст по параграфам
+ original_paragraphs = [p.text for p in sample_document.paragraphs]
+
+ # Проверяем, что все параграфы сохранились
+ result_paragraphs = result_text.split('\n')
+ assert len(result_paragraphs) == len(original_paragraphs)
+
+ # Проверяем, что каждый параграф совпадает с оригиналом
+ for orig, res in zip(original_paragraphs, result_paragraphs):
+ assert orig.strip() == res.strip()
+
+ def test_non_sequential_chunks(self, large_document, doc_entity):
+ """Тест обработки непоследовательных чанков с вставкой многоточий."""
+ strategy = FixedSizeChunkingStrategy(words_per_chunk=10, overlap_words=2)
+
+ # Разбиваем документ
+ entities = strategy.chunk(large_document, doc_entity)
+ chunks = [e for e in entities if e.type == "FixedSizeChunk"]
+
+ # Проверяем, что получили достаточное количество чанков
+ assert len(chunks) >= 5, "Для теста нужно не менее 5 чанков"
+
+ # Отсортируем чанки по индексу
+ sorted_chunks = sorted(chunks, key=lambda c: c.chunk_index or 0)
+
+ # Выберем несколько несмежных чанков (например, 0, 1, 3, 4, 7)
+ selected_indices = [0, 1, 3, 4, 7]
+ selected_chunks = [sorted_chunks[i] for i in selected_indices if i < len(sorted_chunks)]
+
+ # Перемешаем чанки, чтобы убедиться, что сортировка работает
+ import random
+ random.shuffle(selected_chunks)
+
+ # Собираем документ из несмежных чанков
+ result_text = strategy.dechunk(selected_chunks)
+
+ # Проверяем наличие многоточий между непоследовательными чанками
+ assert "\n\n...\n\n" in result_text, "В тексте должно быть многоточие между непоследовательными чанками"
+
+ # Подсчитываем количество многоточий, должно быть 2 группы разрыва (между 1-3 и 4-7)
+ ellipsis_count = result_text.count("\n\n...\n\n")
+ assert ellipsis_count == 2, f"Ожидалось 2 многоточия, получено {ellipsis_count}"
+
+ # Проверяем, что чанки с индексами 0 и 1 идут без многоточия между ними
+ # Для этого находим текст первого чанка и проверяем, что после него нет многоточия
+ first_chunk_text = sorted_chunks[0].text
+ second_chunk_text = sorted_chunks[1].text
+
+ # Проверяем, что текст первого чанка не заканчивается многоточием
+ first_chunk_position = result_text.find(first_chunk_text)
+ second_chunk_position = result_text.find(second_chunk_text, first_chunk_position)
+
+ # Текст между первым и вторым чанком не должен содержать многоточие
+ text_between = result_text[first_chunk_position + len(first_chunk_text):second_chunk_position]
+ assert "\n\n...\n\n" not in text_between, "Не должно быть многоточия между последовательными чанками"
+
+ def test_overlap_addition_in_dechunk(self, large_document, doc_entity):
+ """Тест добавления нахлеста при сборке чанков."""
+ strategy = FixedSizeChunkingStrategy(words_per_chunk=15, overlap_words=5)
+
+ # Разбиваем документ
+ entities = strategy.chunk(large_document, doc_entity)
+ chunks = [e for e in entities if e.type == "FixedSizeChunk"]
+
+ # Отбираем несколько чанков с непустыми overlap_left и overlap_right
+ overlapping_chunks = []
+ for chunk in chunks:
+ if hasattr(chunk, 'overlap_left') and hasattr(chunk, 'overlap_right'):
+ if chunk.overlap_left and chunk.overlap_right:
+ overlapping_chunks.append(chunk)
+ if len(overlapping_chunks) >= 3:
+ break
+
+ # Проверяем, что нашли подходящие чанки
+ assert len(overlapping_chunks) > 0, "Не найдены чанки с нахлестом"
+
+ # Собираем чанки
+ result_text = strategy.dechunk(overlapping_chunks)
+
+ # Проверяем, что нахлесты включены в результат
+ for chunk in overlapping_chunks:
+ if hasattr(chunk, 'overlap_left') and chunk.overlap_left:
+ # Хотя бы часть нахлеста должна присутствовать в тексте
+ # Берем первые три слова нахлеста для проверки
+ overlap_words = chunk.overlap_left.split()[:3]
+ if overlap_words:
+ overlap_sample = " ".join(overlap_words)
+ assert overlap_sample in result_text, f"Левый нахлест не найден в результате: {overlap_sample}"
+
+ if hasattr(chunk, 'overlap_right') and chunk.overlap_right:
+ # Аналогично проверяем правый нахлест
+ overlap_words = chunk.overlap_right.split()[:3]
+ if overlap_words:
+ overlap_sample = " ".join(overlap_words)
+ assert overlap_sample in result_text, f"Правый нахлест не найден в результате: {overlap_sample}"
+
+ # Проверяем обработку предложений
+ for chunk in overlapping_chunks:
+ if hasattr(chunk, 'left_sentence_part') and chunk.left_sentence_part:
+ assert chunk.left_sentence_part in result_text, "Левая часть предложения не найдена в результате"
+
+ if hasattr(chunk, 'right_sentence_part') and chunk.right_sentence_part:
+ assert chunk.right_sentence_part in result_text, "Правая часть предложения не найдена в результате"
\ No newline at end of file
diff --git a/lib/extractor/tests/chunking/test_integration_fixed_size.py b/lib/extractor/tests/chunking/test_integration_fixed_size.py
new file mode 100644
index 0000000000000000000000000000000000000000..119e1d7223d0cf33ae921a3f0a30da79cb687bf2
--- /dev/null
+++ b/lib/extractor/tests/chunking/test_integration_fixed_size.py
@@ -0,0 +1,267 @@
+"""
+Интеграционные тесты для компонентов системы.
+"""
+
+from dataclasses import dataclass
+from unittest.mock import MagicMock
+
+from ntr_fileparser import ParsedDocument, ParsedTextBlock
+
+from ntr_text_fragmentation.core.destructurer import Destructurer
+from ntr_text_fragmentation.core.entity_repository import InMemoryEntityRepository
+from ntr_text_fragmentation.core.injection_builder import InjectionBuilder
+from ntr_text_fragmentation.chunking.specific_strategies.fixed_size import FixedSizeChunk
+
+
+# Создаем простой класс для имитации параграфов документа
+@dataclass
+class MockParagraph:
+ text: str
+
+ def to_string(self) -> str:
+ return self.text
+
+
+class TestFixedSizeChunkingIntegration:
+ """Интеграционные тесты для FixedSizeChunkingStrategy через Destructurer."""
+
+ def test_destructurer_and_injection_builder_integration(self):
+ """
+ Тестирует полный цикл: разбиение документа на чанки и обратную сборку текста.
+ """
+ # Создаем тестовый документ
+ sample_text = (
+ "Это первый параграф тестового документа. Он содержит несколько предложений. "
+ "Эти предложения должны быть корректно обработаны.\n"
+ "Это второй параграф. Он также важен для тестирования.\n"
+ "Это третий параграф, который позволит проверить работу с несколькими блоками текста."
+ )
+
+ paragraphs = [
+ ParsedTextBlock(text=paragraph) for paragraph in sample_text.split('\n')
+ ]
+
+ doc = ParsedDocument(
+ name="Тестовый документ", type="test", paragraphs=paragraphs
+ )
+
+ # Настраиваем Destructurer с параметрами стратегии
+ destructurer = Destructurer(
+ document=doc,
+ strategy_name="fixed_size",
+ words_per_chunk=20, # Небольшой размер для тестирования
+ overlap_words=5, # Небольшой нахлест для тестирования
+ )
+
+ # Получаем сущности из документа
+ entities = destructurer.destructure()
+
+ # Проверяем, что сущности были созданы
+ assert len(entities) > 0
+
+ # Находим документ среди сущностей
+ doc_entity = next((e for e in entities if e.type == "Document"), None)
+ assert doc_entity is not None
+
+ # Находим чанки
+ chunks = [e for e in entities if "Chunk" in e.type]
+ assert len(chunks) > 0
+
+ # Находим связи
+ links = [e for e in entities if e.type == "Link"]
+ assert len(links) > 0
+
+ # Проверяем, что у каждого чанка есть связь с документом
+ for chunk in chunks:
+ assert any(
+ link.target_id == chunk.id and link.source_id == doc_entity.id
+ for link in links
+ if hasattr(link, 'target_id') and hasattr(link, 'source_id')
+ )
+
+ # Создаем репозиторий и сборщик инъекций
+ repository = InMemoryEntityRepository(entities)
+ injection_builder = InjectionBuilder(repository=repository)
+
+ # Получаем идентификаторы чанков для сборки
+ chunk_ids = [chunk.id for chunk in chunks]
+
+ # Собираем текст
+ assembled_text = injection_builder.build(filtered_entities=chunk_ids)
+
+ # Проверяем, что текст был собран
+ assert assembled_text
+
+ # Проверяем наличие ключевых фраз из исходного текста
+ for phrase in ["первый параграф", "второй параграф", "третий параграф"]:
+ assert phrase in assembled_text
+
+ # Проверяем, что порядок параграфов сохранен
+ first_idx = assembled_text.find("первый параграф")
+ second_idx = assembled_text.find("второй параграф")
+ third_idx = assembled_text.find("третий параграф")
+
+ assert 0 <= first_idx < second_idx < third_idx
+
+ def test_add_neighboring_chunks(self):
+ """
+ Тестирует функциональность добавления соседних чанков.
+ """
+ # Создаем тестовый документ с длинным текстом
+ sample_text = "\n".join(
+ [
+ f"Параграф {i}. Этот текст предназначен для тестирования. " * 3
+ for i in range(1, 11)
+ ]
+ )
+
+ paragraphs = [
+ ParsedTextBlock(text=paragraph) for paragraph in sample_text.split('\n')
+ ]
+
+ doc = ParsedDocument(
+ name="Документ для проверки соседей", type="test", paragraphs=paragraphs
+ )
+
+ # Настраиваем Destructurer для создания множества чанков
+ destructurer = Destructurer(
+ document=doc,
+ strategy_name="fixed_size",
+ words_per_chunk=10, # Маленький размер для создания множества чанков
+ overlap_words=2, # Минимальный нахлест
+ )
+
+ # Получаем сущности
+ entities = destructurer.destructure()
+
+ # Находим чанки
+ chunks = [e for e in entities if "Chunk" in e.type]
+ assert len(chunks) > 5 # Должно быть много чанков
+
+ # Берем один чанк из середины
+ middle_chunk = chunks[len(chunks) // 2]
+
+ # Создаем репозиторий и сборщик инъекций
+ repository = InMemoryEntityRepository(entities)
+ injection_builder = InjectionBuilder(repository=repository)
+
+ # Добавляем соседние чанки
+ extended_entities = injection_builder.add_neighboring_chunks(
+ [middle_chunk], max_distance=1
+ )
+
+ # Проверяем, что количество чанков увеличилось, но не равно общему количеству
+ extended_chunks = [e for e in extended_entities if "Chunk" in e.type]
+ assert len(extended_chunks) > 1 # Должно быть больше одного чанка
+ assert len(extended_chunks) < len(chunks) # Но не все чанки
+
+ # Проверяем, что исходный чанк присутствует
+ assert any(chunk.id == middle_chunk.id for chunk in extended_chunks)
+
+ # Проверяем соседние чанки по индексу
+ middle_index = middle_chunk.chunk_index
+ expected_indexes = [middle_index - 1, middle_index, middle_index + 1]
+
+ # Получаем индексы чанков в расширенном наборе
+ extended_indexes = [
+ chunk.chunk_index
+ for chunk in extended_chunks
+ if hasattr(chunk, 'chunk_index')
+ ]
+
+ # Проверяем, что все ожидаемые индексы присутствуют
+ for idx in expected_indexes:
+ if (
+ 0 <= idx < len(chunks)
+ ): # Проверяем, что индекс в пределах допустимых значений
+ assert idx in extended_indexes
+
+ def test_destructurer_fixed_size_basic(self):
+ """
+ Тест для проверки базовой функциональности FixedSizeChunkingStrategy через Destructurer.
+ """
+ # Создаем тестовый документ с одним параграфом
+ test_text = "Это простой тестовый документ для проверки стратегии чанкинга фиксированного размера."
+
+ # Создаем мок ParsedDocument
+ mock_document = MagicMock(spec=ParsedDocument)
+ mock_document.name = "Тестовый документ"
+ mock_document.type = "text"
+
+ # Добавляем параграфы
+ mock_document.paragraphs = [MockParagraph(test_text)]
+
+ # Используем Destructurer с FixedSizeChunkingStrategy
+ destructurer = Destructurer(
+ document=mock_document,
+ strategy_name="fixed_size",
+ words_per_chunk=5, # 5 слов в чанке
+ overlap_words=2, # 2 слова нахлеста
+ )
+
+ # Получаем сущности
+ entities = destructurer.destructure()
+
+ # Проверяем, что сущности были созданы
+ assert len(entities) > 0
+
+ # Находим документ
+ doc_entity = next((e for e in entities if e.type == "Document"), None)
+ assert doc_entity is not None
+
+ # Находим чанки
+ chunks = [e for e in entities if "Chunk" in e.type]
+ assert len(chunks) > 0
+
+ # Проверяем, что у каждого чанка есть текст и индекс
+ for i, chunk in enumerate(sorted(chunks, key=lambda c: c.chunk_index or 0)):
+ assert chunk.text
+ assert chunk.chunk_index == i
+ assert len(chunk.text.split()) <= 5 # не более 5 слов в чанке
+
+ def test_full_cycle_with_builder(self):
+ """
+ Тест полного цикла от Destructurer до InjectionBuilder с произвольными данными.
+ """
+ # Создаем тестовый документ с двумя параграфами
+ paragraph1 = "Первый параграф содержит несколько слов для тестирования."
+ paragraph2 = "Второй параграф также включает в себя некоторое количество слов."
+
+ # Создаем мок ParsedDocument
+ mock_document = MagicMock(spec=ParsedDocument)
+ mock_document.name = "Тестовый документ с параграфами"
+ mock_document.type = "text"
+
+ # Добавляем параграфы
+ mock_document.paragraphs = [
+ MockParagraph(paragraph1),
+ MockParagraph(paragraph2),
+ ]
+
+ # Используем Destructurer с FixedSizeChunkingStrategy
+ destructurer = Destructurer(
+ document=mock_document,
+ strategy_name="fixed_size",
+ words_per_chunk=5, # 5 слов в чанке
+ overlap_words=0, # без нахлеста для простоты тестирования
+ )
+
+ # Получаем сущности
+ entities = destructurer.destructure()
+
+ # Находим чанки
+ chunks = [e for e in entities if "Chunk" in e.type]
+
+ # Создаем репозиторий и сборщик инъекций
+ repository = InMemoryEntityRepository(entities)
+ builder = InjectionBuilder(repository=repository)
+
+ # Собираем документ из всех чанков
+ result = builder.build(filtered_entities=chunks)
+
+ # Проверяем результат
+ assert result # Результат не должен быть пустым
+
+ # Проверяем наличие ключевых слов из обоих параграфов
+ assert "Первый параграф" in result
+ assert "Второй параграф" in result
diff --git a/lib/extractor/tests/conftest.py b/lib/extractor/tests/conftest.py
new file mode 100644
index 0000000000000000000000000000000000000000..52aa6b42907ce102ac3b186f12c9df225fd59b0c
--- /dev/null
+++ b/lib/extractor/tests/conftest.py
@@ -0,0 +1,55 @@
+"""
+Конфигурация pytest для тестов ntr_text_fragmentation.
+"""
+
+from uuid import UUID
+
+import pytest
+
+from ntr_text_fragmentation.models.linker_entity import LinkerEntity
+from tests.custom_entity import CustomEntity # Импортируем наш кастомный класс
+
+
+@pytest.fixture
+def sample_entity():
+ """
+ Фикстура, возвращающая экземпляр LinkerEntity с предустановленными значениями.
+ """
+ return LinkerEntity(
+ id=UUID('12345678-1234-5678-1234-567812345678'),
+ name="Тестовая сущность",
+ text="Текст тестовой сущности",
+ metadata={"test_key": "test_value"}
+ )
+
+
+@pytest.fixture
+def sample_custom_entity():
+ """
+ Фикстура, возвращающая экземпляр CustomEntity с предустановленными значениями.
+ """
+ return CustomEntity(
+ id=UUID('87654321-8765-4321-8765-432187654321'),
+ name="Тестовый кастомный объект",
+ text="Текст кастомного объекта",
+ metadata={"original_key": "original_value"},
+ in_search_text="Текст для поиска кастомного объекта",
+ custom_field1="custom_value",
+ custom_field2=42
+ )
+
+
+@pytest.fixture
+def sample_link():
+ """
+ Фикстура, возвращающая экземпляр LinkerEntity с предустановленными значениями связи.
+ """
+ return LinkerEntity(
+ id=UUID('98765432-9876-5432-9876-543298765432'),
+ name="Тестовая связь",
+ text="Текст тестовой связи",
+ metadata={"test_key": "test_value"},
+ source_id=UUID('12345678-1234-5678-1234-567812345678'),
+ target_id=UUID('87654321-8765-4321-8765-432187654321'),
+ type="Link"
+ )
\ No newline at end of file
diff --git a/lib/extractor/tests/core/__init__.py b/lib/extractor/tests/core/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5dd9750a0b3011610e2263883fd383f84314ebe8
--- /dev/null
+++ b/lib/extractor/tests/core/__init__.py
@@ -0,0 +1,3 @@
+"""
+Модуль с тестами для core-модулей.
+"""
\ No newline at end of file
diff --git a/lib/extractor/tests/core/test_entity_repository.py b/lib/extractor/tests/core/test_entity_repository.py
new file mode 100644
index 0000000000000000000000000000000000000000..3744770f64c0c915718876410393a8fb745eb2ec
--- /dev/null
+++ b/lib/extractor/tests/core/test_entity_repository.py
@@ -0,0 +1,316 @@
+"""
+Тесты для модуля entity_repository.
+"""
+
+from uuid import uuid4
+
+import pytest
+
+from ntr_text_fragmentation.core.entity_repository import \
+ InMemoryEntityRepository
+from ntr_text_fragmentation.models.chunk import Chunk
+from ntr_text_fragmentation.models.document import DocumentAsEntity
+from ntr_text_fragmentation.models.linker_entity import LinkerEntity
+
+
+@pytest.fixture
+def sample_entities():
+ """Создает набор тестовых сущностей."""
+ # Создаем документы как экземпляры DocumentAsEntity вместо LinkerEntity с флагом
+ doc1 = DocumentAsEntity(
+ id=uuid4(),
+ name="document1",
+ text="Документ 1",
+ metadata={"chunking_strategy": "fixed_size"}
+ )
+
+ doc2 = DocumentAsEntity(
+ id=uuid4(),
+ name="document2",
+ text="Документ 2",
+ metadata={"chunking_strategy": "sentence"}
+ )
+
+ # Создаем чанки с индексами
+ chunks_doc1 = [
+ Chunk(id=uuid4(), text=f"Текст чанка {i}", name=f"Чанк {i}", chunk_index=i, metadata={})
+ for i in range(5)
+ ]
+
+ chunks_doc2 = [
+ Chunk(id=uuid4(), text=f"Текст из документа 2, чанк {i}", name=f"Чанк документа 2 - {i}", chunk_index=i, metadata={})
+ for i in range(3)
+ ]
+
+ # Создаем связи между документами и чанками
+ links = []
+ for chunk in chunks_doc1:
+ link = LinkerEntity(
+ id=uuid4(),
+ name="document_to_chunk",
+ text=f"Связь документ-чанк {chunk.name}",
+ metadata={},
+ source_id=doc1.id,
+ target_id=chunk.id,
+ type="Link"
+ )
+ links.append(link)
+
+ for chunk in chunks_doc2:
+ link = LinkerEntity(
+ id=uuid4(),
+ name="document_to_chunk",
+ text=f"Связь документ-чанк {chunk.name}",
+ metadata={},
+ source_id=doc2.id,
+ target_id=chunk.id,
+ type="Link"
+ )
+ links.append(link)
+
+ # Дополнительные связи для тестирования других отношений
+ extra_link = LinkerEntity(
+ id=uuid4(),
+ name="reference",
+ text="Референсная связь",
+ metadata={},
+ source_id=chunks_doc1[0].id,
+ target_id=chunks_doc1[1].id,
+ type="Link"
+ )
+ links.append(extra_link)
+
+ all_entities = [doc1, doc2] + chunks_doc1 + chunks_doc2 + links
+ return {
+ "all": all_entities,
+ "docs": [doc1, doc2],
+ "chunks_doc1": chunks_doc1,
+ "chunks_doc2": chunks_doc2,
+ "links": links
+ }
+
+
+@pytest.fixture
+def repository(sample_entities):
+ """Создает репозиторий с тестовыми данными."""
+ return InMemoryEntityRepository(sample_entities["all"])
+
+
+def test_get_entities_by_ids(repository, sample_entities):
+ """Тест получения сущностей по ID."""
+ # Проверка получения документов
+ doc_ids = [doc.id for doc in sample_entities["docs"]]
+ result = repository.get_entities_by_ids(doc_ids)
+ assert len(result) == 2
+ assert all(entity.id in doc_ids for entity in result)
+
+ # Проверка получения чанков
+ chunk_ids = [chunk.id for chunk in sample_entities["chunks_doc1"]]
+ result = repository.get_entities_by_ids(chunk_ids)
+ assert len(result) == 5
+ assert all(entity.id in chunk_ids for entity in result)
+
+ # Проверка получения несуществующих сущностей
+ non_existent_id = uuid4()
+ result = repository.get_entities_by_ids([non_existent_id])
+ assert len(result) == 0 # Ожидаем пустой список, а не список с None
+
+ # Проверка с пустым списком
+ result = repository.get_entities_by_ids([])
+ assert len(result) == 0
+
+
+def test_get_document_for_chunks(repository, sample_entities):
+ """Тест получения документов для чанков."""
+ # Проверка для чанков из первого документа
+ chunk_ids = [chunk.id for chunk in sample_entities["chunks_doc1"]]
+ result = repository.get_document_for_chunks(chunk_ids)
+ assert len(result) == 1
+ assert result[0].id == sample_entities["docs"][0].id
+
+ # Проверка для чанков из разных документов
+ mixed_chunk_ids = [
+ sample_entities["chunks_doc1"][0].id,
+ sample_entities["chunks_doc2"][0].id
+ ]
+ result = repository.get_document_for_chunks(mixed_chunk_ids)
+ assert len(result) == 2
+ result_ids = [doc.id for doc in result]
+ assert sample_entities["docs"][0].id in result_ids
+ assert sample_entities["docs"][1].id in result_ids
+
+ # Проверка для несуществующих чанков
+ non_existent_id = uuid4()
+ result = repository.get_document_for_chunks([non_existent_id])
+ assert len(result) == 0
+
+ # Проверка с пустым списком
+ result = repository.get_document_for_chunks([])
+ assert len(result) == 0
+
+
+def test_get_neighboring_chunks(repository, sample_entities):
+ """Тест получения соседних чанков."""
+ # Проверка с max_distance=1
+ chunk_ids = [sample_entities["chunks_doc1"][2].id] # Средний чанк (индекс 2)
+ result = repository.get_neighboring_chunks(chunk_ids, max_distance=1)
+ assert len(result) == 2
+ result_indices = [chunk.chunk_index for chunk in result]
+ assert 1 in result_indices # Предыдущий чанк
+ assert 3 in result_indices # Следующий чанк
+
+ # Проверка с max_distance=2
+ result = repository.get_neighboring_chunks(chunk_ids, max_distance=2)
+ assert len(result) == 4
+ result_indices = [chunk.chunk_index for chunk in result]
+ assert all(idx in result_indices for idx in [0, 1, 3, 4])
+
+ # Проверка для граничных чанков (первый чанк)
+ chunk_ids = [sample_entities["chunks_doc1"][0].id] # Первый чанк (индекс 0)
+ result = repository.get_neighboring_chunks(chunk_ids, max_distance=1)
+ assert len(result) == 1
+ assert result[0].chunk_index == 1
+
+ # Проверка для нескольких чанков одновременно
+ chunk_ids = [
+ sample_entities["chunks_doc1"][0].id, # Индекс 0
+ sample_entities["chunks_doc1"][4].id # Индекс 4
+ ]
+ result = repository.get_neighboring_chunks(chunk_ids, max_distance=1)
+ assert len(result) == 2
+ result_indices = [chunk.chunk_index for chunk in result]
+ assert 1 in result_indices # Сосед для индекса 0
+ assert 3 in result_indices # Сосед для индекса 4
+
+ # Проверка с чанками из разных документов
+ chunk_ids = [
+ sample_entities["chunks_doc1"][0].id,
+ sample_entities["chunks_doc2"][0].id
+ ]
+ result = repository.get_neighboring_chunks(chunk_ids, max_distance=1)
+ # Ожидаем чанк с индексом 1 из doc1 и чанк с индексом 1 из doc2
+ assert len(result) == 2
+
+ # Проверка с несуществующими чанками
+ non_existent_id = uuid4()
+ result = repository.get_neighboring_chunks([non_existent_id])
+ assert len(result) == 0
+
+ # Проверка с пустым списком
+ result = repository.get_neighboring_chunks([])
+ assert len(result) == 0
+
+
+def test_get_related_entities(repository, sample_entities):
+ """Тест получения связанных сущностей."""
+ # Получение всех связанных сущностей для документа
+ doc_id = sample_entities["docs"][0].id
+ result = repository.get_related_entities([doc_id])
+ # Ожидаем 5 связей + 5 чанков = 10 сущностей
+ assert len(result) == 10
+
+ # Проверка фильтрации по имени отношения
+ result = repository.get_related_entities([doc_id], relation_name="document_to_chunk")
+ assert len(result) == 10
+
+ result = repository.get_related_entities([doc_id], relation_name="non_existent_relation")
+ assert len(result) == 0
+
+ # Проверка получения связей, где сущность является целью
+ chunk_id = sample_entities["chunks_doc1"][0].id
+ result = repository.get_related_entities([chunk_id], as_target=True)
+ assert len(result) == 2 # 1 связь doc-to-chunk + 1 документ
+
+ # Проверка фильтрации по имени отношения при as_target=True
+ result = repository.get_related_entities(
+ [chunk_id], relation_name="document_to_chunk", as_target=True
+ )
+ assert len(result) == 2
+
+ # Проверка получения связей с несколькими сущностями
+ chunk_ids = [chunk.id for chunk in sample_entities["chunks_doc1"][:2]]
+ result = repository.get_related_entities(chunk_ids, as_target=True)
+ assert len(result) >= 3 # Минимум 2 связи и 1 документ
+
+ # Проверка для несуществующих сущностей
+ non_existent_id = uuid4()
+ result = repository.get_related_entities([non_existent_id])
+ assert len(result) == 0
+
+ # Проверка с пустым списком
+ result = repository.get_related_entities([])
+ assert len(result) == 0
+
+
+def test_add_entities(repository, sample_entities):
+ """Тест добавления сущностей в репозиторий."""
+ # Создаем новые сущности
+ new_doc = DocumentAsEntity(
+ id=uuid4(),
+ name="new_document",
+ text="Новый документ",
+ metadata={"chunking_strategy": "fixed_size"}
+ )
+ new_chunk = Chunk(id=uuid4(), name="Новый чанк", text="Текст нового чанка", chunk_index=0, metadata={})
+ new_link = LinkerEntity(
+ id=uuid4(),
+ name="document_to_chunk",
+ text="Связь новый документ-чанк",
+ metadata={},
+ source_id=new_doc.id,
+ target_id=new_chunk.id,
+ type="Link"
+ )
+
+ # Сохраняем начальное количество сущностей
+ initial_count = len(repository.entities)
+
+ # Добавляем новые сущности
+ repository.add_entities([new_doc, new_chunk, new_link])
+
+ # Проверяем, что сущности добавлены
+ assert len(repository.entities) == initial_count + 3
+
+ # Проверяем, что индексы обновлены
+ assert repository.entities_by_id[new_doc.id] is new_doc
+ assert repository.entities_by_id[new_chunk.id] is new_chunk
+ assert repository.entities_by_id[new_link.id] is new_link
+
+ # Проверяем, что связи обновлены
+ assert new_chunk.id in repository.doc_to_chunks[new_doc.id]
+ assert repository.chunk_to_doc[new_chunk.id] == new_doc.id
+
+ # Проверяем, что можем получить новые сущности
+ result = repository.get_entities_by_ids([new_doc.id, new_chunk.id])
+ assert len(result) == 2
+ assert result[0].id == new_doc.id
+ assert result[1].id == new_chunk.id
+
+ # Проверяем, что можем получить документ для чанка
+ result = repository.get_document_for_chunks([new_chunk.id])
+ assert len(result) == 1
+ assert result[0].id == new_doc.id
+
+
+def test_initialization_with_empty_list():
+ """Тест инициализации репозитория с пустым списком сущностей."""
+ repository = InMemoryEntityRepository([])
+ assert len(repository.entities) == 0
+ assert len(repository.entities_by_id) == 0
+ assert len(repository.chunks) == 0
+ assert len(repository.docs) == 0
+
+ # Проверка методов с пустым репозиторием
+ assert repository.get_entities_by_ids([uuid4()]) == []
+ assert repository.get_document_for_chunks([uuid4()]) == []
+ assert repository.get_neighboring_chunks([uuid4()]) == []
+ assert repository.get_related_entities([uuid4()]) == []
+
+
+def test_initialization_without_arguments():
+ """Тест инициализации репозитория без аргументов."""
+ repository = InMemoryEntityRepository()
+ assert len(repository.entities) == 0
+ assert len(repository.entities_by_id) == 0
+ assert len(repository.chunks) == 0
+ assert len(repository.docs) == 0
\ No newline at end of file
diff --git a/lib/extractor/tests/custom_entity.py b/lib/extractor/tests/custom_entity.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca96e042f843e368436f54c1aef5297b94ee6d75
--- /dev/null
+++ b/lib/extractor/tests/custom_entity.py
@@ -0,0 +1,106 @@
+from uuid import UUID
+
+from ntr_text_fragmentation.models.linker_entity import (LinkerEntity,
+ register_entity)
+
+
+@register_entity
+class CustomEntity(LinkerEntity):
+ """Пользовательский класс-наследник LinkerEntity для тестирования сериализации и десериализации."""
+
+ def __init__(
+ self,
+ id: UUID,
+ name: str,
+ text: str,
+ metadata: dict,
+ custom_field1: str,
+ custom_field2: int,
+ in_search_text: str | None = None,
+ source_id: UUID | None = None,
+ target_id: UUID | None = None,
+ number_in_relation: int | None = None,
+ type: str = "CustomEntity"
+ ):
+ super().__init__(
+ id=id,
+ name=name,
+ text=text,
+ metadata=metadata,
+ in_search_text=in_search_text,
+ source_id=source_id,
+ target_id=target_id,
+ number_in_relation=number_in_relation,
+ type=type
+ )
+ self.custom_field1 = custom_field1
+ self.custom_field2 = custom_field2
+
+ def deserialize(self, entity: LinkerEntity) -> 'CustomEntity':
+ """Реализация метода десериализации для кастомного класса."""
+ custom_field1 = entity.metadata.get('_custom_field1', '')
+ custom_field2 = entity.metadata.get('_custom_field2', 0)
+
+ # Создаем чистые метаданные без служебных полей
+ clean_metadata = {k: v for k, v in entity.metadata.items()
+ if not k.startswith('_')}
+
+ return CustomEntity(
+ id=entity.id,
+ name=entity.name,
+ text=entity.text,
+ in_search_text=entity.in_search_text,
+ metadata=clean_metadata,
+ source_id=entity.source_id,
+ target_id=entity.target_id,
+ number_in_relation=entity.number_in_relation,
+ custom_field1=custom_field1,
+ custom_field2=custom_field2
+ )
+
+ @classmethod
+ def deserialize(cls, entity: LinkerEntity) -> 'CustomEntity':
+ """
+ Классовый метод для десериализации.
+ Необходим для работы с реестром классов.
+
+ Args:
+ entity: Сериализованная сущность
+
+ Returns:
+ Десериализованный экземпляр CustomEntity
+ """
+ custom_field1 = entity.metadata.get('_custom_field1', '')
+ custom_field2 = entity.metadata.get('_custom_field2', 0)
+
+ # Создаем чистые метаданные без служебных полей
+ clean_metadata = {k: v for k, v in entity.metadata.items()
+ if not k.startswith('_')}
+
+ return CustomEntity(
+ id=entity.id,
+ name=entity.name,
+ text=entity.text,
+ in_search_text=entity.in_search_text,
+ metadata=clean_metadata,
+ source_id=entity.source_id,
+ target_id=entity.target_id,
+ number_in_relation=entity.number_in_relation,
+ custom_field1=custom_field1,
+ custom_field2=custom_field2
+ )
+
+ def __eq__(self, other):
+ """Переопределяем метод сравнения для проверки равенства объектов."""
+ if not isinstance(other, CustomEntity):
+ return False
+
+ # Используем базовое сравнение из LinkerEntity, которое уже учитывает поля связи
+ base_equality = super().__eq__(other)
+
+ # Дополнительно проверяем кастомные поля
+ return (
+ base_equality
+ and self.custom_field1 == other.custom_field1
+ and self.custom_field2 == other.custom_field2
+ )
\ No newline at end of file
diff --git a/lib/extractor/tests/models/__init__.py b/lib/extractor/tests/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3fbd3962a30308926d4f36f574a5e32e56681f46
--- /dev/null
+++ b/lib/extractor/tests/models/__init__.py
@@ -0,0 +1,3 @@
+"""
+Тесты для моделей данных.
+"""
\ No newline at end of file
diff --git a/lib/extractor/tests/models/test_entity.py b/lib/extractor/tests/models/test_entity.py
new file mode 100644
index 0000000000000000000000000000000000000000..23fba4f7996e695bbb5b29d90eb96f3983d78414
--- /dev/null
+++ b/lib/extractor/tests/models/test_entity.py
@@ -0,0 +1,314 @@
+import uuid
+from uuid import UUID
+
+from ntr_text_fragmentation.models.linker_entity import LinkerEntity
+from tests.custom_entity import CustomEntity
+
+
+class TestEntity:
+ """Набор тестов для проверки класса LinkerEntity."""
+
+ def test_create_entity(self):
+ """Тест создания объекта LinkerEntity с заданными параметрами."""
+ entity_id = uuid.uuid4()
+ name = "Тестовая сущность"
+ text = "Это текст тестовой сущности"
+ in_search_text = "Текст для поиска"
+ metadata = {"key": "value"}
+
+ entity = LinkerEntity(
+ id=entity_id,
+ name=name,
+ text=text,
+ in_search_text=in_search_text,
+ metadata=metadata
+ )
+
+ assert entity.id == entity_id
+ assert entity.name == name
+ assert entity.text == text
+ assert entity.in_search_text == in_search_text
+ assert entity.metadata == metadata
+ assert entity.type == "Entity" # Проверка значения по умолчанию
+ assert entity.source_id is None # Проверка опциональных полей
+ assert entity.target_id is None
+ assert entity.number_in_relation is None
+ assert entity.is_link() is False # Не является связью
+
+ def test_create_entity_with_link_fields(self):
+ """Тест создания объекта LinkerEntity с полями связи."""
+ entity_id = uuid.uuid4()
+ source_id = uuid.uuid4()
+ target_id = uuid.uuid4()
+
+ entity = LinkerEntity(
+ id=entity_id,
+ name="Тестовая связывающая сущность",
+ text="Это текст связывающей сущности",
+ in_search_text="Текст для поиска",
+ metadata={"key": "value"},
+ source_id=source_id,
+ target_id=target_id,
+ number_in_relation=1
+ )
+
+ assert entity.id == entity_id
+ assert entity.source_id == source_id
+ assert entity.target_id == target_id
+ assert entity.number_in_relation == 1
+ assert entity.is_link() is True # Является связью
+
+ def test_entity_fixture(self, sample_entity):
+ """Тест использования фикстуры с предустановленными значениями."""
+ assert sample_entity.id == UUID('12345678-1234-5678-1234-567812345678')
+ assert sample_entity.name == "Тестовая сущность"
+ assert sample_entity.text == "Текст тестовой сущности"
+ assert sample_entity.in_search_text is None # По умолчанию None
+ assert sample_entity.metadata == {"test_key": "test_value"}
+ assert sample_entity.type == "Entity"
+ assert sample_entity.is_link() is False
+
+ def test_auto_id_generation(self):
+ """Тест автоматической генерации ID, если он не указан."""
+ entity = LinkerEntity(
+ id=None,
+ name="Тест",
+ text="Текст",
+ metadata={}
+ )
+
+ assert entity.id is not None
+ assert isinstance(entity.id, UUID)
+
+ def test_invalid_link_fields(self):
+ """Тест создания сущности с некорректными полями связи."""
+ # Пробуем создать сущность только с source_id
+ try:
+ LinkerEntity(
+ id=uuid.uuid4(),
+ name="Некорректная сущность",
+ text="Текст некорректной сущности",
+ metadata={},
+ source_id=uuid.uuid4()
+ )
+ assert False, "Создание сущности только с source_id должно вызывать исключение"
+ except ValueError:
+ pass
+
+ # Пробуем создать сущность только с target_id
+ try:
+ LinkerEntity(
+ id=uuid.uuid4(),
+ name="Некорректная сущность",
+ text="Текст некорректной сущности",
+ metadata={},
+ target_id=uuid.uuid4()
+ )
+ assert False, "Создание сущности только с target_id должно вызывать исключение"
+ except ValueError:
+ pass
+
+ def test_custom_type(self):
+ """Тест использования пользовательского типа."""
+ custom_type = "CustomEntity"
+ entity = LinkerEntity(
+ id=uuid.uuid4(),
+ name="Тест",
+ text="Текст",
+ metadata={},
+ type=custom_type
+ )
+
+ assert entity.type == custom_type
+
+ def test_to_string(self):
+ """Тест метода __str__."""
+ # Тест стандартного вывода
+ entity = LinkerEntity(
+ id=uuid.uuid4(),
+ name="Тест",
+ text="Текст",
+ metadata={}
+ )
+ expected_string = "Тест: Текст"
+ assert str(entity) == expected_string
+
+ # Тест с использованием in_search_text
+ entity_with_search = LinkerEntity(
+ id=uuid.uuid4(),
+ name="Тест",
+ text="Текст",
+ in_search_text="Текст для поиска",
+ metadata={}
+ )
+ assert str(entity_with_search) == "Текст для поиска"
+
+ def test_compare_same_entities(self):
+ """Тест сравнения одинаковых сущностей."""
+ entity_id = uuid.uuid4()
+ name = "Сущность"
+ text = "Текст"
+ entity_type = "TestEntity"
+
+ entity1 = LinkerEntity(
+ id=entity_id,
+ name=name,
+ text=text,
+ in_search_text="Текст для поиска 1",
+ metadata={"a": 1},
+ type=entity_type
+ )
+
+ entity2 = LinkerEntity(
+ id=entity_id,
+ name=name,
+ text=text,
+ in_search_text="Текст для поиска 2", # Этот параметр не учитывается в сравнении
+ metadata={"b": 2}, # Этот параметр не учитывается в сравнении
+ type=entity_type
+ )
+
+ assert entity1 == entity2
+
+ def test_compare_different_entities(self, sample_entity):
+ """Тест сравнения разных сущностей."""
+ # Проверка с другим ID
+ different_id_entity = LinkerEntity(
+ id=uuid.uuid4(),
+ name=sample_entity.name,
+ text=sample_entity.text,
+ metadata=sample_entity.metadata
+ )
+ assert sample_entity != different_id_entity
+
+ # Проверка с другим именем
+ different_name_entity = LinkerEntity(
+ id=sample_entity.id,
+ name="Другое имя",
+ text=sample_entity.text,
+ metadata=sample_entity.metadata
+ )
+ assert sample_entity != different_name_entity
+
+ # Проверка с другим текстом
+ different_text_entity = LinkerEntity(
+ id=sample_entity.id,
+ name=sample_entity.name,
+ text="Другой текст",
+ metadata=sample_entity.metadata
+ )
+ assert sample_entity != different_text_entity
+
+ # Проверка с другим типом
+ different_type_entity = LinkerEntity(
+ id=sample_entity.id,
+ name=sample_entity.name,
+ text=sample_entity.text,
+ metadata=sample_entity.metadata,
+ type="ДругойТип"
+ )
+ assert sample_entity != different_type_entity
+
+ def test_compare_with_other_class(self, sample_entity):
+ """Тест сравнения с объектом другого класса."""
+ # Создаем объект другого класса
+ class OtherClass:
+ pass
+
+ other = OtherClass()
+ assert sample_entity != other
+
+ def test_serialize(self, sample_custom_entity):
+ """Тест метода serialize для кастомного класса-наследника Entity."""
+ # Сериализуем объект
+ serialized = sample_custom_entity.serialize()
+
+ # Проверяем, что сериализованный объект - это базовый Entity
+ assert isinstance(serialized, LinkerEntity)
+ assert serialized.id == sample_custom_entity.id
+ assert serialized.name == "Тестовый кастомный объект"
+ assert serialized.text == "Текст кастомного объекта"
+ # Проверяем, что тип соответствует имени класса согласно новой логике
+ assert serialized.type == "CustomEntity"
+
+ # Проверяем, что кастомные поля автоматически сохранены в метаданных
+ assert "_custom_field1" in serialized.metadata
+ assert "_custom_field2" in serialized.metadata
+ assert serialized.metadata["_custom_field1"] == "custom_value"
+ assert serialized.metadata["_custom_field2"] == 42
+ assert serialized.metadata["original_key"] == "original_value"
+
+ def test_deserialize(self):
+ """Тест метода deserialize для кастомного класса-наследника Entity."""
+ # Создаем базовый Entity с метаданными, как будто это сериализованный CustomEntity
+ entity_id = uuid.uuid4()
+ serialized_entity = LinkerEntity(
+ id=entity_id,
+ name="Тестовый кастомный объект",
+ text="Текст кастомного объекта",
+ in_search_text="Текст для поиска",
+ metadata={
+ "_custom_field1": "custom_value",
+ "_custom_field2": 42,
+ "original_key": "original_value"
+ },
+ type="CustomEntity" # Используем имя класса при сериализации
+ )
+
+ # Десериализуем объект
+ template = CustomEntity(
+ id=uuid.uuid4(), # Не важно, будет заменено
+ name="",
+ text="",
+ in_search_text=None,
+ metadata={},
+ custom_field1="",
+ custom_field2=0
+ )
+ deserialized = template.deserialize(serialized_entity)
+
+ # Проверяем, что десериализованный объект корректно восстановил все поля
+ assert isinstance(deserialized, CustomEntity)
+ assert deserialized.id == entity_id
+ assert deserialized.name == "Тестовый кастомный объект"
+ assert deserialized.text == "Текст кастомного объекта"
+ assert deserialized.in_search_text == "Текст для поиска"
+ assert deserialized.custom_field1 == "custom_value"
+ assert deserialized.custom_field2 == 42
+ assert deserialized.metadata == {"original_key": "original_value"}
+
+ def test_serialize_deserialize_roundtrip(self, sample_custom_entity):
+ """Тест полного цикла сериализации и десериализации."""
+ # Полный цикл сериализация -> десериализация
+ serialized = sample_custom_entity.serialize()
+ deserialized = sample_custom_entity.deserialize(serialized)
+
+ # Проверяем, что объект после полного цикла идентичен исходному
+ assert deserialized == sample_custom_entity
+ assert deserialized.id == sample_custom_entity.id
+ assert deserialized.name == sample_custom_entity.name
+ assert deserialized.text == sample_custom_entity.text
+ assert deserialized.in_search_text == sample_custom_entity.in_search_text
+ assert deserialized.metadata == sample_custom_entity.metadata
+ assert deserialized.type == sample_custom_entity.type
+ assert deserialized.custom_field1 == sample_custom_entity.custom_field1
+ assert deserialized.custom_field2 == sample_custom_entity.custom_field2
+
+ def test_static_deserialize_method(self, sample_custom_entity):
+ """Тест статического метода deserialize класса LinkerEntity."""
+ # CustomEntity уже зарегистрирован через декоратор @register_entity
+
+ # Сериализуем объект
+ serialized = sample_custom_entity.serialize()
+
+ # Десериализуем через статический метод LinkerEntity.deserialize
+ deserialized = LinkerEntity.deserialize(serialized)
+
+ # Проверяем, что десериализация прошла правильно
+ assert isinstance(deserialized, CustomEntity)
+ assert deserialized == sample_custom_entity
+ assert deserialized.id == sample_custom_entity.id
+ assert deserialized.name == sample_custom_entity.name
+ assert deserialized.text == sample_custom_entity.text
+ assert deserialized.custom_field1 == sample_custom_entity.custom_field1
+ assert deserialized.custom_field2 == sample_custom_entity.custom_field2
\ No newline at end of file
diff --git a/lib/extractor/tests/models/test_link.py b/lib/extractor/tests/models/test_link.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d48df25f1e3bc2bee0caa500fd2fb01d6e5d28a
--- /dev/null
+++ b/lib/extractor/tests/models/test_link.py
@@ -0,0 +1,207 @@
+import uuid
+from uuid import UUID
+
+from ntr_text_fragmentation.models.linker_entity import LinkerEntity
+
+
+class TestLink:
+ """Набор тестов для проверки функциональности связей (линков) с использованием LinkerEntity."""
+
+ def test_create_link(self):
+ """Тест создания объекта LinkerEntity с параметрами связи."""
+ link_id = uuid.uuid4()
+ name = "Тестовая связь"
+ text = "Это текст тестовой связи"
+ in_search_text = "Текст для поиска связи"
+ metadata = {"key": "value"}
+ source_id = uuid.uuid4()
+ target_id = uuid.uuid4()
+
+ link = LinkerEntity(
+ id=link_id,
+ name=name,
+ text=text,
+ in_search_text=in_search_text,
+ metadata=metadata,
+ source_id=source_id,
+ target_id=target_id,
+ type="Link"
+ )
+
+ assert link.id == link_id
+ assert link.name == name
+ assert link.text == text
+ assert link.in_search_text == in_search_text
+ assert link.metadata == metadata
+ assert link.source_id == source_id
+ assert link.target_id == target_id
+ assert link.type == "Link"
+ assert link.number_in_relation is None
+ assert link.is_link() is True
+
+ def test_link_fixture(self, sample_link):
+ """Тест использования фикстуры с предустановленными значениями."""
+ assert sample_link.id == UUID('98765432-9876-5432-9876-543298765432')
+ assert sample_link.name == "Тестовая связь"
+ assert sample_link.text == "Текст тестовой связи"
+ assert sample_link.in_search_text is None # Значение по умолчанию
+ assert sample_link.metadata == {"test_key": "test_value"}
+ assert sample_link.source_id == UUID('12345678-1234-5678-1234-567812345678')
+ assert sample_link.target_id == UUID('87654321-8765-4321-8765-432187654321')
+ assert sample_link.type == "Link"
+ assert sample_link.number_in_relation is None
+ assert sample_link.is_link() is True
+
+ def test_link_with_number_in_relation(self):
+ """Тест создания объекта LinkerEntity с указанным номером в связи."""
+ link_id = uuid.uuid4()
+ name = "Тестовая связь"
+ text = "Это текст тестовой связи"
+ in_search_text = "Текст для поиска связи"
+ metadata = {"key": "value"}
+ source_id = uuid.uuid4()
+ target_id = uuid.uuid4()
+ number_in_relation = 1
+
+ link = LinkerEntity(
+ id=link_id,
+ name=name,
+ text=text,
+ in_search_text=in_search_text,
+ metadata=metadata,
+ source_id=source_id,
+ target_id=target_id,
+ number_in_relation=number_in_relation,
+ type="Link"
+ )
+
+ assert link.id == link_id
+ assert link.name == name
+ assert link.text == text
+ assert link.in_search_text == in_search_text
+ assert link.metadata == metadata
+ assert link.source_id == source_id
+ assert link.target_id == target_id
+ assert link.type == "Link"
+ assert link.number_in_relation == number_in_relation
+ assert link.is_link() is True
+
+ def test_link_to_string(self):
+ """Тест метода __str__()."""
+ # Без in_search_text
+ link = LinkerEntity(
+ id=uuid.uuid4(),
+ name="Тестовая связь",
+ text="Текст тестовой связи",
+ metadata={},
+ source_id=uuid.uuid4(),
+ target_id=uuid.uuid4(),
+ type="Link"
+ )
+ expected_string = "Тестовая связь: Текст тестовой связи"
+ assert str(link) == expected_string
+
+ # С in_search_text
+ link_with_search_text = LinkerEntity(
+ id=uuid.uuid4(),
+ name="Тестовая связь",
+ text="Текст тестовой связи",
+ in_search_text="Специальный текст для поиска",
+ metadata={},
+ source_id=uuid.uuid4(),
+ target_id=uuid.uuid4(),
+ type="Link"
+ )
+ assert str(link_with_search_text) == "Специальный текст для поиска"
+
+ def test_link_compare(self, sample_link):
+ """Тест метода __eq__ для сравнения связей."""
+ # Создаем копию с теми же основными атрибутами
+ same_link = LinkerEntity(
+ id=sample_link.id,
+ name=sample_link.name,
+ text=sample_link.text,
+ in_search_text="Другой текст для поиска", # Отличается от sample_link
+ metadata={}, # Отличается от sample_link
+ source_id=sample_link.source_id,
+ target_id=sample_link.target_id,
+ type="Link"
+ )
+
+ # Должны быть равны по __eq__, так как метаданные и in_search_text не учитываются в сравнении
+ assert sample_link == same_link
+
+ # Создаем связь с другим ID
+ different_id_link = LinkerEntity(
+ id=uuid.uuid4(),
+ name=sample_link.name,
+ text=sample_link.text,
+ metadata=sample_link.metadata,
+ source_id=sample_link.source_id,
+ target_id=sample_link.target_id,
+ type="Link"
+ )
+
+ assert sample_link != different_id_link
+
+ def test_invalid_link_creation(self):
+ """Тест создания некорректной связи (без source_id или target_id)."""
+ link_id = uuid.uuid4()
+
+ # Пробуем создать связь только с source_id
+ try:
+ LinkerEntity(
+ id=link_id,
+ name="Некорректная связь",
+ text="Текст некорректной связи",
+ metadata={},
+ source_id=uuid.uuid4()
+ )
+ assert False, "Создание связи только с source_id должно вызывать исключение"
+ except ValueError:
+ pass
+
+ # Пробуем создать связь только с target_id
+ try:
+ LinkerEntity(
+ id=link_id,
+ name="Некорректная связь",
+ text="Текст некорректной связи",
+ metadata={},
+ target_id=uuid.uuid4()
+ )
+ assert False, "Создание связи только с target_id должно вызывать исключение"
+ except ValueError:
+ pass
+
+ def test_linker_entity_as_link(self):
+ """Тест использования LinkerEntity как связи."""
+ entity_id = uuid.uuid4()
+ source_id = uuid.uuid4()
+ target_id = uuid.uuid4()
+
+ # Создаем LinkerEntity с source_id и target_id
+ linking_entity = LinkerEntity(
+ id=entity_id,
+ name="Связывающая сущность",
+ text="Текст связывающей сущности",
+ metadata={},
+ source_id=source_id,
+ target_id=target_id
+ )
+
+ # Проверяем, что LinkerEntity может выступать как связь
+ assert linking_entity.is_link() is True
+ assert linking_entity.source_id == source_id
+ assert linking_entity.target_id == target_id
+
+ # Создаем обычную сущность
+ regular_entity = LinkerEntity(
+ id=uuid.uuid4(),
+ name="Обычная сущность",
+ text="Текст обычной сущности",
+ metadata={}
+ )
+
+ # Проверяем, что обычная сущность не является связью
+ assert regular_entity.is_link() is False
\ No newline at end of file
diff --git a/lib/parser/.cursor/rules/relative-imports.mdc b/lib/parser/.cursor/rules/relative-imports.mdc
new file mode 100644
index 0000000000000000000000000000000000000000..3a25625e3c9df79f548679224554b28dad67e258
--- /dev/null
+++ b/lib/parser/.cursor/rules/relative-imports.mdc
@@ -0,0 +1,6 @@
+---
+description: USE ONLY RELATIVE IMPORTS
+globs:
+alwaysApply: false
+---
+Use only relative imports
\ No newline at end of file
diff --git a/lib/parser/.gitignore b/lib/parser/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..5b468bf10ea128552e5f71e88981c3d5c207bff5
--- /dev/null
+++ b/lib/parser/.gitignore
@@ -0,0 +1,13 @@
+use_it/*
+test_output/
+test_input/
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+*.pyw
+*.pyz
+
+*.egg-info/
+*.dist-info/
+*.build-info/
diff --git a/lib/parser/README.md b/lib/parser/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..c79ac76dab4c469499f46581deb6f4f3384374f2
--- /dev/null
+++ b/lib/parser/README.md
@@ -0,0 +1,109 @@
+# Универсальный парсер документов
+
+Библиотека `ntr_fileparser` предоставляет универсальный интерфейс для извлечения структурированных данных из документов различных форматов.
+
+## Возможности
+
+- Извлечение текстовых блоков, таблиц, изображений и формул из документов
+- Поддержка различных форматов файлов:
+ - XML (реализовано)
+ - DOCX (реализовано)
+ - PDF (в разработке)
+ - HTML (в разработке)
+ - Markdown (в разработке)
+ - Email (в разработке)
+ - DOC (в разработке)
+- Единый интерфейс для работы со всеми форматами
+- Расширяемая архитектура для добавления новых парсеров
+
+## Установка
+
+```bash
+pip install git+ssh://git@gitlab.ntrlab.ru:textai/parsers/parser.git
+```
+
+## Использование
+
+### Базовый пример
+
+```python
+from ntr_fileparser import UniversalParser
+
+# Создание экземпляра парсера
+parser = UniversalParser()
+
+# Парсинг документа по пути к файлу
+document = parser.parse_by_path("path/to/document.xml")
+
+# Доступ к данным
+print(f"Название документа: {document.name}")
+print(f"Количество параграфов: {len(document.paragraphs)}")
+print(f"Количество таблиц: {len(document.tables)}")
+```
+
+### Работа с содержимым файла
+
+```python
+from ntr_fileparser import UniversalParser, FileType
+
+# Создание экземпляра парсера
+parser = UniversalParser()
+
+# Открытие файла
+with open("path/to/document.xml", "rb") as f:
+ # Парсинг содержимого файла с указанием типа файла
+ # Тип файла можно указать как объект FileType или как строку расширения
+ document = parser.parse(f, FileType.XML)
+ # Или
+ document = parser.parse(f, ".xml")
+
+ if document:
+ # Работа с документом
+ for paragraph in document.paragraphs:
+ print(paragraph.text)
+```
+
+### Применение функций к документу
+
+```python
+from ntr_fileparser import UniversalParser
+
+# Создание экземпляра парсера
+parser = UniversalParser()
+
+# Парсинг документа
+document = parser.parse_by_path("path/to/document.xml")
+
+# Применение функции ко всем текстовым элементам документа
+document.apply(lambda text: text.upper())
+```
+
+## Архитектура
+
+Система построена на основе следующих компонентов:
+
+1. **UniversalParser** - основной класс, предоставляющий единый интерфейс для работы со всеми форматами документов.
+2. **AbstractParser** - абстрактный класс парсера, определяющий интерфейс для всех конкретных парсеров.
+3. **ParserFactory** - фабрика парсеров, отвечающая за выбор подходящего парсера для конкретного документа.
+4. **FileType** - перечисление поддерживаемых типов файлов.
+5. **Специализированные парсеры** - классы для работы с конкретными форматами документов (XMLParser, PDFParser и т.д.).
+6. **Структуры данных** - классы для представления структуры документа (ParsedDocument, ParsedTextBlock, ParsedTable и т.д.).
+
+### Диаграмма архитектуры
+
+
+
+## Требования
+
+- Python 3.11+
+- Зависимости автоматически устанавливаются при установке библиотеки
+
+## Разработка
+
+Для разработки и тестирования:
+
+```bash
+git clone git@gitlab.ntrlab.ru:textai/parsers/parser.git
+cd parser
+pip install -e .
+```
diff --git a/lib/parser/__init__.py b/lib/parser/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b417dfd2d6d7c71a56bf3b9b6c4b27e1781c6359
--- /dev/null
+++ b/lib/parser/__init__.py
@@ -0,0 +1,26 @@
+"""
+Пакет для парсинга документов различных форматов.
+
+Этот пакет предоставляет универсальный парсер документов и необходимые структуры данных
+для представления содержимого документа.
+
+Данный файл __init__.py используется только для импорта пакетов при использовании
+git-submodule. В библиотечном варианте точкой входа является файл parsing/__init__.py
+"""
+
+from .ntr_fileparser import *
+
+__all__ = [
+ 'FileType',
+ 'UniversalParser',
+ 'ParsedDocument',
+ 'ParsedMeta',
+ 'ParsedStructure',
+ 'ParsedTextBlock',
+ 'ParsedTable',
+ 'ParsedSubtable',
+ 'ParsedRow',
+ 'ParsedImage',
+ 'ParsedFormula',
+ 'TableTag',
+]
diff --git a/lib/parser/docs/NTR_FileParser.svg b/lib/parser/docs/NTR_FileParser.svg
new file mode 100644
index 0000000000000000000000000000000000000000..bdb59d0b7879c34c8322fdf9db5bcd65ac8e5ba5
--- /dev/null
+++ b/lib/parser/docs/NTR_FileParser.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/__init__.py b/lib/parser/ntr_fileparser/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..090cfad886b08a9df4bd757def9512eb93387f81
--- /dev/null
+++ b/lib/parser/ntr_fileparser/__init__.py
@@ -0,0 +1,35 @@
+"""
+Модуль парсинга документов различных форматов.
+
+Данный файл является точкой входа при использовании библиотеки через pip.
+"""
+
+from .data_classes import (
+ ParsedDocument,
+ ParsedFormula,
+ ParsedImage,
+ ParsedMeta,
+ ParsedRow,
+ ParsedStructure,
+ ParsedSubtable,
+ ParsedTable,
+ ParsedTextBlock,
+ TableTag,
+)
+from .parsers.universal_parser import UniversalParser
+from .parsers.file_types import FileType
+
+__all__ = [
+ 'FileType',
+ 'UniversalParser',
+ 'ParsedDocument',
+ 'ParsedMeta',
+ 'ParsedStructure',
+ 'ParsedTextBlock',
+ 'ParsedTable',
+ 'ParsedSubtable',
+ 'ParsedRow',
+ 'ParsedImage',
+ 'ParsedFormula',
+ 'TableTag',
+]
diff --git a/lib/parser/ntr_fileparser/data_classes/__init__.py b/lib/parser/ntr_fileparser/data_classes/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d1a2d4cb324f0cfe5e757dfdb1e27e63c4884be
--- /dev/null
+++ b/lib/parser/ntr_fileparser/data_classes/__init__.py
@@ -0,0 +1,25 @@
+"""
+Модуль содержит датаклассы для представления структуры документа.
+"""
+
+from .parsed_document import ParsedDocument
+from .parsed_formula import ParsedFormula
+from .parsed_image import ParsedImage
+from .parsed_meta import ParsedMeta
+from .parsed_structure import ParsedStructure
+from .parsed_table import ParsedRow, ParsedSubtable, ParsedTable, TableTag
+from .parsed_text_block import ParsedTextBlock, TextStyle
+
+__all__ = [
+ 'ParsedStructure',
+ 'ParsedMeta',
+ 'ParsedTextBlock',
+ 'TableTag',
+ 'ParsedRow',
+ 'ParsedSubtable',
+ 'ParsedTable',
+ 'ParsedImage',
+ 'ParsedFormula',
+ 'ParsedDocument',
+ 'TextStyle',
+]
diff --git a/lib/parser/ntr_fileparser/data_classes/parsed_document.py b/lib/parser/ntr_fileparser/data_classes/parsed_document.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc836718e3193ed8423ade61db514d3ad7105453
--- /dev/null
+++ b/lib/parser/ntr_fileparser/data_classes/parsed_document.py
@@ -0,0 +1,105 @@
+"""
+Модуль содержит класс для представления структуры документа.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+from .parsed_formula import ParsedFormula
+from .parsed_image import ParsedImage
+from .parsed_meta import ParsedMeta
+from .parsed_structure import ParsedStructure
+from .parsed_table import ParsedTable
+from .parsed_text_block import ParsedTextBlock
+
+
+@dataclass
+class ParsedDocument(ParsedStructure):
+ """
+ Документ, полученный в результате парсинга.
+ """
+ name: str = ""
+ type: str = ""
+ meta: ParsedMeta = field(default_factory=ParsedMeta)
+ paragraphs: list[ParsedTextBlock] = field(default_factory=list)
+ tables: list[ParsedTable] = field(default_factory=list)
+ images: list[ParsedImage] = field(default_factory=list)
+ formulas: list[ParsedFormula] = field(default_factory=list)
+
+ def to_string(self) -> str:
+ """
+ Преобразует документ в строковое представление.
+
+ Returns:
+ str: Строковое представление документа.
+ """
+ result = [f"Документ: {self.name} (тип: {self.type})"]
+
+ if self.paragraphs:
+ result.append("\nПараграфы:")
+ for p in self.paragraphs:
+ result.append(p.to_string())
+
+ if self.tables:
+ result.append("\nТаблицы:")
+ for t in self.tables:
+ result.append(t.to_string())
+
+ if self.images:
+ result.append("\nИзображения:")
+ for i in self.images:
+ result.append(i.to_string())
+
+ if self.formulas:
+ result.append("\nФормулы:")
+ for f in self.formulas:
+ result.append(f.to_string())
+
+ return "\n".join(result)
+
+ def apply(self, func: Callable[[str], str]) -> None:
+ """
+ Применяет функцию ко всем строковым элементам документа.
+
+ Args:
+ func (Callable[[str], str]): Функция для применения к текстовым элементам.
+ """
+ self.name = func(self.name)
+ self.type = func(self.type)
+
+ # Применяем к параграфам
+ for p in self.paragraphs:
+ p.apply(func)
+
+ # Применяем к таблицам
+ for t in self.tables:
+ t.apply(func)
+
+ # Применяем к изображениям
+ for i in self.images:
+ i.apply(func)
+
+ # Применяем к формулам
+ for f in self.formulas:
+ f.apply(func)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует документ в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление документа.
+ """
+ # Преобразуем тип в строку, если это объект FileType
+ type_str = str(self.type) if not isinstance(self.type, str) else self.type
+
+ result = {
+ 'name': self.name,
+ 'type': type_str,
+ 'meta': self.meta.to_dict(),
+ 'paragraphs': [p.to_dict() for p in self.paragraphs],
+ 'tables': [t.to_dict() for t in self.tables],
+ 'images': [i.to_dict() for i in self.images],
+ 'formulas': [f.to_dict() for f in self.formulas]
+ }
+ return result
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/data_classes/parsed_formula.py b/lib/parser/ntr_fileparser/data_classes/parsed_formula.py
new file mode 100644
index 0000000000000000000000000000000000000000..9be46232d74d8655485c501b10e5fee41a9c9d94
--- /dev/null
+++ b/lib/parser/ntr_fileparser/data_classes/parsed_formula.py
@@ -0,0 +1,63 @@
+"""
+Модуль содержит класс для представления формул в документе.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from .parsed_structure import DocumentElement
+
+
+@dataclass
+class ParsedFormula(DocumentElement):
+ """
+ Формула из документа.
+ """
+
+ title: str | None = None
+ latex: str = ""
+ # Номер формулы в документе
+ formula_number: str | None = None
+ # Дополнительное описание/пояснение формулы
+ description: str | None = None
+
+ def to_string(self) -> str:
+ """
+ Преобразует формулу в строковое представление.
+
+ Returns:
+ str: Строковое представление формулы.
+ """
+ title_str = f"{self.title}: " if self.title else ""
+ return f"{title_str}Формула: {self.latex}"
+
+ def apply(self, func: Callable[[str], str]) -> None:
+ """
+ Применяет функцию к текстовым элементам формулы.
+
+ Args:
+ func (Callable[[str], str]): Функция для применения к текстовым элементам.
+ """
+ if self.title:
+ self.title = func(self.title)
+ self.latex = func(self.latex)
+ if self.description:
+ self.description = func(self.description)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует формулу в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление формулы.
+ """
+ result = {
+ 'title': self.title,
+ 'latex': self.latex,
+ 'formula_number': self.formula_number,
+ 'description': self.description,
+ 'page_number': self.page_number,
+ 'index_in_document': self.index_in_document,
+ 'referenced_element_index': self.referenced_element_index
+ }
+ return result
diff --git a/lib/parser/ntr_fileparser/data_classes/parsed_image.py b/lib/parser/ntr_fileparser/data_classes/parsed_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..2507643fd893cedd6f7370cfd6de27aa9db091a7
--- /dev/null
+++ b/lib/parser/ntr_fileparser/data_classes/parsed_image.py
@@ -0,0 +1,61 @@
+"""
+Модуль содержит класс для представления изображений в документе.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+from .parsed_structure import DocumentElement
+
+
+@dataclass
+class ParsedImage(DocumentElement):
+ """
+ Изображение из документа (нереализованный класс).
+ """
+
+ title: str | None = None
+ image: bytes = field(default_factory=bytes)
+ # Размеры изображения (в пикселях или единицах документа)
+ width: int | None = None
+ height: int | None = None
+ # Подпись/описание изображения
+ caption: str | None = None
+
+ def to_string(self) -> str:
+ """
+ Преобразует информацию об изображении в строковое представление.
+
+ Returns:
+ str: Строковое представление информации об изображении.
+ """
+ return f"Изображение: {self.title if self.title else 'Без названия'}"
+
+ def apply(self, func: Callable[[str], str]) -> None:
+ """
+ Применяет функцию к текстовым элементам изображения.
+
+ Args:
+ func (Callable[[str], str]): Функция для применения к текстовым элементам.
+ """
+ if self.title:
+ self.title = func(self.title)
+ if self.caption:
+ self.caption = func(self.caption)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует изображение в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление изображения.
+ """
+ return {
+ 'title': self.title,
+ 'width': self.width,
+ 'height': self.height,
+ 'caption': self.caption,
+ 'page_number': self.page_number,
+ 'index_in_document': self.index_in_document,
+ 'referenced_element_index': self.referenced_element_index
+ }
diff --git a/lib/parser/ntr_fileparser/data_classes/parsed_meta.py b/lib/parser/ntr_fileparser/data_classes/parsed_meta.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab7bc1e3a2d658ccf4d739ea3f2d0b4e6562ac54
--- /dev/null
+++ b/lib/parser/ntr_fileparser/data_classes/parsed_meta.py
@@ -0,0 +1,39 @@
+"""
+Модуль содержит класс для метаданных документа.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+
+@dataclass
+class ParsedMeta:
+ """
+ Метаданные документа.
+ """
+ date: datetime | str = field(default_factory=datetime.now)
+ owner: str | None = None
+ source: str | None = None
+ status: str | None = None
+ note: dict | None = None
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует метаданные в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление метаданных.
+ """
+ date_value = self.date
+ # Конвертируем datetime в строку, если это объект datetime
+ if isinstance(self.date, datetime):
+ date_value = self.date.isoformat()
+
+ return {
+ 'date': date_value,
+ 'owner': self.owner,
+ 'source': self.source,
+ 'status': self.status,
+ 'note': self.note,
+ }
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/data_classes/parsed_structure.py b/lib/parser/ntr_fileparser/data_classes/parsed_structure.py
new file mode 100644
index 0000000000000000000000000000000000000000..13303bc8fb7cd4c45b32d92d4ae3bd6c11e44349
--- /dev/null
+++ b/lib/parser/ntr_fileparser/data_classes/parsed_structure.py
@@ -0,0 +1,58 @@
+"""
+Модуль содержит базовый интерфейс для всех структурных элементов документа.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Callable
+
+
+@dataclass
+class ParsedStructure(ABC):
+ """
+ Базовый абстрактный класс для всех структурных элементов документа.
+ """
+
+ @abstractmethod
+ def to_string(self) -> str:
+ """
+ Преобразует структуру в строковое представление.
+
+ Returns:
+ str: Строковое представление структуры.
+ """
+ pass
+
+ @abstractmethod
+ def apply(self, func: Callable[[str], str]) -> None:
+ """
+ Применяет трансформации к строковым элементам,
+ аналогично функции apply в pandas.
+
+ Args:
+ func (Callable[[str], str]): Функция для применения к строковым элементам.
+ """
+ pass
+
+ @abstractmethod
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует структуру в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление структуры.
+ """
+ pass
+
+
+@dataclass
+class DocumentElement(ParsedStructure):
+ """
+ Базовый класс для всех элементов документа (параграфы, таблицы, изображения, формулы).
+ """
+
+ # Номер страницы, на которой находится элемент
+ page_number: int | None = None
+
+ # Индекс элемента в документе (порядковый номер)
+ index_in_document: int | None = None
diff --git a/lib/parser/ntr_fileparser/data_classes/parsed_table.py b/lib/parser/ntr_fileparser/data_classes/parsed_table.py
new file mode 100644
index 0000000000000000000000000000000000000000..819adc43841d56779404f2a4da7ca18e6439634c
--- /dev/null
+++ b/lib/parser/ntr_fileparser/data_classes/parsed_table.py
@@ -0,0 +1,530 @@
+"""
+Модуль содержит классы для представления таблиц в документе.
+"""
+
+import warnings
+from dataclasses import asdict, dataclass, field
+from typing import Any, Callable, Optional
+
+from .parsed_structure import DocumentElement
+from .parsed_text_block import TextStyle
+
+
+@dataclass
+class TableTag:
+ """
+ Тег для классификации таблицы.
+ """
+
+ name: str = ""
+ value: str = ""
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует тег таблицы в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление тега таблицы.
+ """
+ return asdict(self)
+
+
+@dataclass
+class ParsedRow(DocumentElement):
+ """
+ Строка таблицы.
+ """
+
+ index: int | str | None = None
+ cells: list[str] = field(default_factory=list)
+ style: TextStyle = field(default_factory=TextStyle)
+ anchors: list[str] = field(default_factory=list)
+ links: list[str] = field(default_factory=list)
+ is_header: bool = False
+
+ def to_string(
+ self,
+ header: Optional['ParsedRow'] = None,
+ note: Optional[str] = None,
+ ) -> str:
+ """
+ Преобразует строку таблицы в строковое представление в виде маркированного списка.
+
+ Args:
+ header (Optional[ParsedRow]): Заголовок столбцов для форматирования.
+ note (Optional[str]): Примечание к строке. Не будет использовано, если строка не содержит *,
+ а в примечании есть *
+
+ Returns:
+ str: Строковое представление строки таблицы.
+ """
+ if not self.cells:
+ return ""
+
+ # Если у нас есть хедер, то форматируем как "ключ: значение" с маркерами
+ if header:
+ if len(header.cells) != len(self.cells):
+ raise ValueError("Количество ячеек в строке и хедере не совпадает")
+ result = '\n'.join(
+ f"- {header.cells[i]}: {self.cells[i]}" for i in range(len(self.cells))
+ )
+
+ # Если у нас есть только две колонки, форматируем как "ключ: значение" с маркером
+ elif len(self.cells) == 2:
+ result = f"- {self.cells[0].strip()}: {self.cells[1].strip()}"
+
+ # Иначе просто форматируем все ячейки через разделитель
+ else:
+ result = '\n'.join(f"- {cell.strip()}" for cell in self.cells)
+
+ if note:
+
+ if ('*' in result) != ('*' in note):
+ return result
+ else:
+ return f"{result}\nПримечание: {note}"
+
+ return result
+
+ def apply(self, func: Callable[[str], str]) -> None:
+ """
+ Применяет функцию ко всем ячейкам строки.
+
+ Args:
+ func (Callable[[str], str]): Функция для применения к текстовым элементам.
+ """
+ self.cells = [func(cell) for cell in self.cells]
+ self.anchors = [func(anchor) for anchor in self.anchors]
+ self.links = [func(link) for link in self.links]
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует строку таблицы в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление строки таблицы.
+ """
+ return {
+ 'index': self.index,
+ 'cells': self.cells,
+ 'style': self.style.to_dict(),
+ 'anchors': self.anchors,
+ 'links': self.links,
+ 'is_header': self.is_header,
+ 'page_number': self.page_number,
+ 'index_in_document': self.index_in_document,
+ }
+
+
+@dataclass
+class ParsedSubtable(DocumentElement):
+ """
+ Подтаблица внутри основной таблицы.
+ """
+
+ title: str | None = None
+ header: ParsedRow | None = None
+ rows: list[ParsedRow] = field(default_factory=list)
+
+ def to_string(
+ self,
+ header: Optional['ParsedRow'] = None,
+ note: Optional[str] = None,
+ ) -> str:
+ """
+ Преобразует подтаблицу в строковое представление.
+
+ Returns:
+ str: Строковое представление подтаблицы.
+ """
+ if self.header:
+ header = self.header
+
+ result = []
+ if self.title:
+ result.append(f"## {self.title}")
+
+ if len(self.rows) == 0:
+ if header:
+ result.append(header.to_string(note=note))
+ if note:
+ result.append(f"Примечание: {note}")
+
+ # Обрабатываем каждую строку таблицы
+ for i, row in enumerate(self.rows, start=1):
+ # Добавляем номер строки (начиная с 1)
+ result.append(f"### Строка {i}")
+ result.append(row.to_string(header=header, note=note))
+
+ return "\n".join(result)
+
+ def apply(self, func: Callable[[str], str]) -> None:
+ """
+ Применяет функцию ко всем элементам подтаблицы.
+
+ Args:
+ func (Callable[[str], str]): Функция для применения к текстовым элементам.
+ """
+ if self.title:
+ self.title = func(self.title)
+ if self.header:
+ self.header.apply(func)
+ for row in self.rows:
+ row.apply(func)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует подтаблицу в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление подтаблицы.
+ """
+ result = {'title': self.title, 'rows': [row.to_dict() for row in self.rows]}
+ if self.header:
+ result['header'] = self.header.to_dict()
+
+ # Добавляем поля из DocumentElement
+ result['page_number'] = self.page_number
+ result['index_in_document'] = self.index_in_document
+
+ return result
+
+ def has_merged_cells(self) -> bool:
+ """
+ Проверяет наличие объединенных ячеек в подтаблице.
+
+ Returns:
+ bool: True, если в подтаблице есть строки с разным количеством ячеек.
+ """
+ if not self.rows:
+ return False
+
+ # Получаем количество ячеек в строках
+ cell_counts = [len(row.cells) for row in self.rows]
+ if len(set(cell_counts)) > 1:
+ return True
+
+ return False
+
+
+@dataclass
+class ParsedTable(DocumentElement):
+ """
+ Таблица из документа.
+ """
+
+ title: str | None = None
+ note: str | None = None
+ classified_tags: list[TableTag] = field(default_factory=list)
+ index: list[str] = field(default_factory=list)
+ headers: list[ParsedRow] = field(default_factory=list)
+ subtables: list[ParsedSubtable] = field(default_factory=list)
+ table_style: dict[str, Any] = field(default_factory=dict)
+ title_index_in_paragraphs: int | None = None
+
+ def to_string(self) -> str:
+ """
+ Преобразует таблицу в строковое представление.
+
+ Returns:
+ str: Строковое представление таблицы.
+ """
+ # Формируем заголовок таблицы
+ table_header = ""
+ if self.title:
+ table_header = f"# {self.title}"
+
+ final_result = []
+
+ common_header = None
+ if self.headers:
+ common_header = ParsedRow(
+ cells=[
+ '/'.join(header.cells[i] for header in self.headers)
+ for i in range(len(self.headers[0].cells))
+ ]
+ )
+
+ if len(self.subtables) == 0:
+ if common_header:
+ final_result.append(common_header.to_string(note=self.note))
+ else:
+ final_result.append(self.note)
+
+ # Обрабатываем каждую подтаблицу
+ for subtable in self.subtables:
+ # Получаем строковое представление подтаблицы
+ subtable_lines = subtable.to_string(common_header, self.note).split('\n')
+
+ # Для каждой линии в подтаблице
+ current_block = []
+ for line in subtable_lines:
+ # Если это начало новой строки (заголовок строки)
+ if line.startswith('### Строка'):
+ # Если у нас уже есть блок данных, добавляем его с дополнительным переносом
+ if current_block:
+ final_result.append('\n'.join(current_block))
+ final_result.append("") # Дополнительный перенос между строками
+
+ # Начинаем новый блок с заголовка таблицы
+ current_block = []
+ if table_header:
+ current_block.append(table_header)
+
+ # Если у подтаблицы есть заголовок, добавляем его
+ if subtable.title:
+ current_block.append(f"## {subtable.title}")
+
+ # Добавляем заголовок строки
+ current_block.append(line)
+ else:
+ # Добавляем данные строки
+ current_block.append(line)
+
+ # Добавляем последний блок, если он есть
+ if current_block:
+ final_result.append('\n'.join(current_block))
+ final_result.append("") # Дополнительный перенос между блоками
+
+ return '\n'.join(final_result)
+
+ def apply(self, func: Callable[[str], str]) -> None:
+ """
+ Применяет функцию ко всем элементам таблицы.
+
+ Args:
+ func (Callable[[str], str]): Функция для применения к текстовым элементам.
+ """
+ if self.title:
+ self.title = func(self.title)
+ self.note = func(self.note)
+ self.index = [func(idx) for idx in self.index]
+ for tag in self.classified_tags:
+ tag.name = func(tag.name)
+ tag.value = func(tag.value)
+ for header in self.headers:
+ header.apply(func)
+ for subtable in self.subtables:
+ subtable.apply(func)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует таблицу в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление таблицы.
+ """
+ result = {
+ 'title': self.title,
+ 'note': self.note,
+ 'classified_tags': [tag.to_dict() for tag in self.classified_tags],
+ 'index': self.index,
+ 'headers': [header.to_dict() for header in self.headers],
+ 'subtables': [subtable.to_dict() for subtable in self.subtables],
+ 'table_style': self.table_style,
+ 'page_number': self.page_number,
+ 'index_in_document': self.index_in_document,
+ 'title_index_in_paragraphs': self.title_index_in_paragraphs,
+ }
+ return result
+
+ def has_merged_cells(self) -> bool:
+ """
+ Проверяет наличие объединенных ячеек в таблице.
+
+ Returns:
+ bool: True, если в таблице есть строки с разным количеством ячеек.
+ """
+ # Проверяем заголовки
+ if self.headers:
+ header_cell_counts = [len(header.cells) for header in self.headers]
+ if len(set(header_cell_counts)) > 1:
+ return True
+
+ expected_cell_count = header_cell_counts[0] if header_cell_counts else 0
+ else:
+ expected_cell_count = 0
+
+ # Проверяем подтаблицы
+ for subtable in self.subtables:
+ if subtable.has_merged_cells():
+ return True
+
+ # Проверяем соответствие количества ячеек заголовку
+ if subtable.rows and expected_cell_count > 0:
+ for row in subtable.rows:
+ if len(row.cells) != expected_cell_count:
+ return True
+
+ return False
+
+ def to_pandas(self, merged_ok: bool = False) -> Optional['pandas.DataFrame']: # type: ignore
+ """
+ Преобразует таблицу в pandas DataFrame.
+
+ Args:
+ merged_ok (bool): Флаг, указывающий, допустимы ли объединенные ячейки.
+ Если False и обнаружены объединенные ячейки, будет выдано предупреждение.
+
+ Returns:
+ pandas.DataFrame: DataFrame, представляющий таблицу.
+
+ Примечание:
+ Этот метод требует установленного пакета pandas.
+ """
+ try:
+ import pandas as pd
+ except ImportError:
+ raise ImportError(
+ "Для использования to_pandas требуется установить pandas."
+ )
+
+ # Проверка объединенных ячеек
+ if not merged_ok and self.has_merged_cells():
+ warnings.warn(
+ "Таблица содержит объединенные ячейки, что может привести к некорректному "
+ "отображению в DataFrame. Установите параметр merged_ok=True, чтобы скрыть это предупреждение."
+ )
+
+ # Собираем данные для DataFrame
+ data = []
+
+ # Заголовки столбцов
+ columns = []
+ if self.headers:
+ # Объединяем многострочные заголовки, используя разделитель '->'
+ if len(self.headers) > 1:
+ # Собираем все строки заголовков
+ header_cells = []
+ for i in range(len(self.headers[0].cells)):
+ header_values = [
+ header.cells[i] if i < len(header.cells) else ""
+ for header in self.headers
+ ]
+ header_cells.append(" -> ".join(filter(None, header_values)))
+ columns = header_cells
+ else:
+ columns = self.headers[0].cells
+
+ # Собираем данные из подтаблиц
+ for subtable in self.subtables:
+ # Если есть заголовок подтаблицы, добавляем его как строку с пустыми значениями
+ if subtable.title:
+ row_data = (
+ [subtable.title] + [""] * (len(columns) - 1)
+ if columns
+ else [subtable.title]
+ )
+ data.append(row_data)
+
+ # Добавляем данные из строк подтаблицы
+ for row in subtable.rows:
+ row_data = row.cells
+
+ # Если количество ячеек не совпадает с количеством столбцов, заполняем пустыми
+ if columns and len(row_data) < len(columns):
+ row_data.extend([""] * (len(columns) - len(row_data)))
+
+ data.append(row_data)
+
+ # Создаем DataFrame
+ if not columns:
+ # Если нет заголовков, определяем максимальное количество столбцов
+ max_cols = max([len(row) for row in data]) if data else 0
+ df = pd.DataFrame(data)
+ else:
+ df = pd.DataFrame(data, columns=columns)
+
+ # Добавляем название таблицы как атрибут
+ if self.title:
+ df.attrs['title'] = self.title
+
+ # Добавляем примечание как атрибут
+ if self.note:
+ df.attrs['note'] = self.note
+
+ return df
+
+ def to_markdown(self, merged_ok: bool = False) -> str:
+ """
+ Преобразует таблицу в формат Markdown.
+
+ Args:
+ merged_ok (bool): Флаг, указывающий, допустимы ли объединенные ячейки.
+ Если False и обнаружены объединенные ячейки, будет выдано предупреждение.
+
+ Returns:
+ str: Markdown представление таблицы.
+ """
+ # Проверка объединенных ячеек
+ if not merged_ok and self.has_merged_cells():
+ warnings.warn(
+ "Таблица содержит объединенные ячейки, что может привести к некорректному "
+ "отображению в Markdown. Установите параметр merged_ok=True, чтобы скрыть это предупреждение."
+ )
+
+ lines = []
+
+ # Добавляем заголовок таблицы, если он есть
+ if self.title:
+ lines.append(f"**{self.title}**\n")
+
+ # Если есть заголовок таблицы, используем его
+ if self.headers:
+ # Берем первую строку заголовка
+ header_cells = self.headers[0].cells
+
+ # Формируем строку заголовка
+ header_line = "| " + " | ".join(header_cells) + " |"
+ lines.append(header_line)
+
+ # Формируем разделительную строку
+ separator_line = "| " + " | ".join(["---"] * len(header_cells)) + " |"
+ lines.append(separator_line)
+
+ # Если есть дополнительные строки заголовка, добавляем их
+ for i in range(1, len(self.headers)):
+ subheader_cells = self.headers[i].cells
+ if len(subheader_cells) < len(header_cells):
+ subheader_cells.extend(
+ [""] * (len(header_cells) - len(subheader_cells))
+ )
+ subheader_line = (
+ "| " + " | ".join(subheader_cells[: len(header_cells)]) + " |"
+ )
+ lines.append(subheader_line)
+
+ # Обходим подтаблицы
+ for subtable in self.subtables:
+ # Если есть заголовок подтаблицы, добавляем его как строку
+ if subtable.title:
+ lines.append(
+ f"| **{subtable.title}** | "
+ + " | ".join([""] * (len(header_cells) - 1))
+ + " |"
+ )
+
+ # Добавляем строки подтаблицы
+ for row in subtable.rows:
+ row_cells = row.cells
+
+ # Если количество ячеек не совпадает с количеством заголовков, добавляем пустые
+ if len(row_cells) < len(header_cells):
+ row_cells.extend([""] * (len(header_cells) - len(row_cells)))
+
+ row_line = "| " + " | ".join(row_cells[: len(header_cells)]) + " |"
+ lines.append(row_line)
+ else:
+ # Если заголовка нет, просто выводим строки как текст
+ for subtable in self.subtables:
+ if subtable.title:
+ lines.append(f"**{subtable.title}**")
+
+ for row in subtable.rows:
+ lines.append(row.to_string())
+
+ # Добавляем примечание, если оно есть
+ if self.note:
+ lines.append(f"\n*Примечание: {self.note}*")
+
+ return "\n".join(lines)
diff --git a/lib/parser/ntr_fileparser/data_classes/parsed_text_block.py b/lib/parser/ntr_fileparser/data_classes/parsed_text_block.py
new file mode 100644
index 0000000000000000000000000000000000000000..4683384835e5c047b38f1346d6dad1ac3e710d0a
--- /dev/null
+++ b/lib/parser/ntr_fileparser/data_classes/parsed_text_block.py
@@ -0,0 +1,116 @@
+"""
+Модуль содержит класс для представления текстовых блоков документа.
+"""
+
+from dataclasses import asdict, dataclass, field
+from typing import Any, Callable
+
+from .parsed_structure import DocumentElement
+
+
+@dataclass
+class TextStyle:
+ """
+ Стиль текстового блока.
+ """
+ # Стиль параграфа
+ paragraph_style: str = ""
+ paragraph_style_name: str = ""
+ alignment: str = ""
+
+ # Автонумерация
+ has_numbering: bool = False
+ numbering_level: int = 0
+ numbering_id: str = ""
+ numbering_format: str = "" # decimal, bullet, roman, etc.
+
+ # Флаги для bold
+ fully_bold: bool = False
+ partly_bold: bool = False
+
+ # Флаги для italic
+ fully_italic: bool = False
+ partly_italic: bool = False
+
+ # Флаги для underline
+ fully_underlined: bool = False
+ partly_underlined: bool = False
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует стиль в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление стиля.
+ """
+ return asdict(self)
+
+
+@dataclass
+class ParsedTextBlock(DocumentElement):
+ """
+ Текстовый блок документа.
+ """
+
+ text: str = ""
+ style: TextStyle = field(default_factory=TextStyle)
+ anchors: list[str] = field(default_factory=list) # Список идентификаторов якорей (закладок)
+ links: list[str] = field(default_factory=list) # Список идентификаторов ссылок
+
+ # Технические метаданные о блоке
+ metadata: list[dict[str, Any]] = field(default_factory=list) # Для хранения технической информации
+
+ # Примечания и сноски к тексту
+ footnotes: list[dict[str, Any]] = field(default_factory=list) # Для хранения сносок
+
+ title_of_table: int | None = None
+
+
+ def to_string(self) -> str:
+ """
+ Преобразует текстовый блок в строковое представление.
+
+ Returns:
+ str: Текст блока.
+ """
+ return self.text
+
+ def apply(self, func: Callable[[str], str]) -> None:
+ """
+ Применяет функцию к тексту блока.
+
+ Args:
+ func (Callable[[str], str]): Функция для применения к тексту.
+ """
+ self.text = func(self.text)
+ # Применяем к текстовым значениям в метаданных
+ if isinstance(self.metadata, list):
+ for item in self.metadata:
+ if isinstance(item, dict) and 'text' in item:
+ item['text'] = func(item['text'])
+
+ # Применяем к сноскам
+ if isinstance(self.footnotes, list):
+ for note in self.footnotes:
+ if isinstance(note, dict) and 'text' in note:
+ note['text'] = func(note['text'])
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Преобразует текстовый блок в словарь.
+
+ Returns:
+ dict[str, Any]: Словарное представление текстового блока.
+ """
+ result = {
+ 'text': self.text,
+ 'style': self.style.to_dict(),
+ 'anchors': self.anchors,
+ 'links': self.links,
+ 'metadata': self.metadata,
+ 'footnotes': self.footnotes,
+ 'page_number': self.page_number,
+ 'index_in_document': self.index_in_document,
+ 'title_of_table': self.title_of_table,
+ }
+ return result
diff --git a/lib/parser/ntr_fileparser/parsers/__init__.py b/lib/parser/ntr_fileparser/parsers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c711f892ceb7f92fa7a4b7c64009fd82fee0050a
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/__init__.py
@@ -0,0 +1,22 @@
+"""
+Модуль содержит парсеры для различных форматов документов.
+"""
+
+from .abstract_parser import AbstractParser
+from .parser_factory import ParserFactory
+from .specific_parsers import (DocParser, DocxParser, EmailParser, HTMLParser,
+ MarkdownParser, PDFParser, XMLParser)
+from .universal_parser import UniversalParser
+
+__all__ = [
+ 'AbstractParser',
+ 'ParserFactory',
+ 'UniversalParser',
+ 'DocParser',
+ 'DocxParser',
+ 'EmailParser',
+ 'HTMLParser',
+ 'MarkdownParser',
+ 'PDFParser',
+ 'XMLParser'
+]
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/abstract_parser.py b/lib/parser/ntr_fileparser/parsers/abstract_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..7e2360e30db3b407afdbb36944720e58e8a9eb35
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/abstract_parser.py
@@ -0,0 +1,142 @@
+"""
+Модуль с абстрактным классом парсера.
+"""
+
+import os
+from abc import ABC, abstractmethod
+from typing import BinaryIO
+
+from ..data_classes import ParsedDocument
+from .file_types import FileType
+
+
+class AbstractParser(ABC):
+ """
+ Абстрактный класс парсера документов.
+
+ Все конкретные парсеры должны наследоваться от этого класса
+ и реализовывать метод parse.
+ """
+
+ def __init__(self, file_types: FileType | list[FileType] | str | list[str] | None = None):
+ """
+ Инициализирует парсер.
+
+ Args:
+ file_types: Поддерживаемые типы файлов. Может быть одним из:
+ - FileType - объект перечисления
+ - list[FileType] - список объектов перечисления
+ - str - строка с расширением файла (с точкой, например ".xml")
+ - list[str] - список строк с расширениями
+ - None - если не указан, парсер не ограничен типами
+ """
+ self.file_types = []
+
+ if file_types is None:
+ return
+
+ # Преобразуем одиночный FileType в список
+ if isinstance(file_types, FileType):
+ self.file_types = [file_types]
+ # Преобразуем список FileType в список
+ elif isinstance(file_types, list) and all(isinstance(ft, FileType) for ft in file_types):
+ self.file_types = file_types
+ # Преобразуем строку расширения в FileType
+ elif isinstance(file_types, str):
+ try:
+ self.file_types = [FileType.from_extension(file_types)]
+ except ValueError:
+ # Если не удалось найти подходящий FileType, создаем пустой список
+ self.file_types = []
+ # Преобразуем список строк расширений в список FileType
+ elif isinstance(file_types, list) and all(isinstance(ft, str) for ft in file_types):
+ self.file_types = []
+ for ext in file_types:
+ try:
+ self.file_types.append(FileType.from_extension(ext))
+ except ValueError:
+ pass
+
+ def _supported_extension(self, ext: str) -> bool:
+ """
+ Проверяет, поддерживается ли расширение файла.
+
+ Этот метод должен быть переопределен в наследниках
+ для указания поддерживаемых расширений.
+
+ Args:
+ ext (str): Расширение файла с точкой (.pdf, .docx и т.д.).
+
+ Returns:
+ bool: True, если расширение поддерживается, иначе False.
+ """
+ if not self.file_types:
+ try:
+ FileType.from_extension(ext)
+ return True
+ except ValueError:
+ return False
+
+ ext = ext.lower()
+ for file_type in self.file_types:
+ for supported_ext in file_type.value:
+ if ext == supported_ext.lower():
+ return True
+ return False
+
+ def supports_file(self, file: str | BinaryIO | FileType) -> bool:
+ """
+ Проверяет, может ли парсер обработать файл.
+
+ Args:
+ file: Может быть одним из:
+ - str: Путь к файлу
+ - BinaryIO: Объект файла
+ - FileType: Конкретный тип файла
+
+ Returns:
+ bool: True, если парсер поддерживает файл, иначе False.
+ """
+ # Если передан FileType, проверяем его наличие в списке поддерживаемых
+ if isinstance(file, FileType):
+ return file in self.file_types
+
+ # Если переданы пустые file_types и не строка, не можем определить тип
+ if not self.file_types and not isinstance(file, str):
+ return False
+
+ # Если передан путь к файлу, проверяем расширение
+ if isinstance(file, str):
+ _, ext = os.path.splitext(file)
+ return self._supported_extension(ext)
+
+ # Если передан бинарный объект, считаем что подходит
+ # (конкретный тип будет проверен при вызове parse)
+ return True
+
+ @abstractmethod
+ def parse(self, file: BinaryIO, file_type: FileType | None = None) -> ParsedDocument:
+ """
+ Парсит документ из объекта файла и возвращает его структурное представление.
+
+ Args:
+ file (BinaryIO): Объект файла для парсинга.
+ file_type (FileType | None): Тип файла, если известен.
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+ """
+ pass
+
+ @abstractmethod
+ def parse_by_path(self, file_path: str) -> ParsedDocument:
+ """
+ Парсит документ по пути к файлу и возвращает его структурное представление.
+
+ Args:
+ file_path (str): Путь к файлу для парсинга.
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+ """
+ pass
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/file_types.py b/lib/parser/ntr_fileparser/parsers/file_types.py
new file mode 100644
index 0000000000000000000000000000000000000000..0dbd7762abdd8c36586f723d60977809f72301bc
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/file_types.py
@@ -0,0 +1,51 @@
+"""
+Модуль с перечислением поддерживаемых типов файлов.
+"""
+
+from enum import Enum
+
+
+class FileType(Enum):
+ """
+ Перечисление поддерживаемых типов файлов.
+ """
+ XML = [".xml"]
+ DOCX = [".docx"]
+ DOC = [".doc"]
+ PDF = [".pdf"]
+ HTML = [".html", ".htm"]
+ MD = [".md", ".markdown"]
+ EML = [".eml"]
+
+ @classmethod
+ def from_extension(cls, ext: str) -> 'FileType':
+ """
+ Получает тип файла по расширению.
+
+ Args:
+ ext (str): Расширение файла (с точкой).
+
+ Returns:
+ FileType: Тип файла.
+
+ Raises:
+ ValueError: Если расширение не поддерживается.
+ """
+ ext = ext.lower()
+ for file_type in cls:
+ if ext in [e.lower() for e in file_type.value]:
+ return file_type
+ raise ValueError(f"Unsupported file extension: {ext}")
+
+ @classmethod
+ def get_supported_extensions(cls) -> list[str]:
+ """
+ Возвращает список всех поддерживаемых расширений.
+
+ Returns:
+ list[str]: Список расширений с точкой.
+ """
+ extensions = []
+ for file_type in cls:
+ extensions.extend(file_type.value)
+ return extensions
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/parser_factory.py b/lib/parser/ntr_fileparser/parsers/parser_factory.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f9c803713c6043af1401a4d2e029b200d8765bc
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/parser_factory.py
@@ -0,0 +1,54 @@
+"""
+Модуль с фабрикой парсеров для различных форматов документов.
+"""
+
+import logging
+
+from .abstract_parser import AbstractParser
+from .file_types import FileType
+
+logger = logging.getLogger(__name__)
+
+
+class ParserFactory:
+ """
+ Фабрика парсеров документов.
+
+ Отвечает за выбор подходящего парсера для конкретного документа.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует фабрику парсеров.
+ """
+ self.parsers: list[AbstractParser] = []
+
+ def register_parser(self, parser: AbstractParser) -> None:
+ """
+ Регистрирует парсер в фабрике.
+
+ Args:
+ parser (AbstractParser): Парсер для регистрации.
+ """
+ self.parsers.append(parser)
+ logger.debug(f"Зарегистрирован парсер: {parser.__class__.__name__}")
+
+ def get_parser(self, file: str | FileType) -> AbstractParser | None:
+ """
+ Возвращает подходящий парсер для файла.
+
+ Args:
+ file: Может быть одним из:
+ - str: Путь к файлу для определения подходящего парсера.
+ - FileType: Тип файла для определения подходящего парсера.
+
+ Returns:
+ AbstractParser | None: Подходящий парсер или None, если такой не найден.
+ """
+ for parser in self.parsers:
+ if parser.supports_file(file):
+ logger.debug(f"Выбран парсер {parser.__class__.__name__} для файла {file}")
+ return parser
+
+ logger.warning(f"Не найден подходящий парсер для {file}")
+ return None
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/__init__.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1aff70cc293ce73d25738e5d04064a3599ce981
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/__init__.py
@@ -0,0 +1,21 @@
+"""
+Модуль, содержащий конкретные реализации парсеров для различных форматов документов.
+"""
+
+from .doc_parser import DocParser
+from .docx_parser import DocxParser
+from .email_parser import EmailParser
+from .html_parser import HTMLParser
+from .markdown_parser import MarkdownParser
+from .pdf_parser import PDFParser
+from .xml_parser import XMLParser
+
+__all__ = [
+ 'DocParser',
+ 'DocxParser',
+ 'EmailParser',
+ 'HTMLParser',
+ 'MarkdownParser',
+ 'PDFParser',
+ 'XMLParser'
+]
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/doc_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/doc_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a455e86207a1524ad7d73ad2d8909787a59ad7e
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/doc_parser.py
@@ -0,0 +1,91 @@
+"""
+Модуль с парсером для DOC документов.
+"""
+
+import logging
+import os
+from typing import BinaryIO
+
+from ...data_classes import ParsedDocument
+from ..abstract_parser import AbstractParser
+from ..file_types import FileType
+
+logger = logging.getLogger(__name__)
+
+
+class DocParser(AbstractParser):
+ """
+ Парсер для старых документов Word формата DOC.
+
+ Примечание: На данный момент реализация является заглушкой.
+ В будущем будет использоваться библиотека textract или pywin32.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует DOC парсер.
+ """
+ super().__init__(FileType.DOC)
+
+ def parse_by_path(self, file_path: str) -> ParsedDocument:
+ """
+ Парсит DOC документ по пути к файлу и возвращает его структурное представление.
+
+ Args:
+ file_path (str): Путь к DOC файлу для парсинга.
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если файл не существует или не может быть прочитан.
+ NotImplementedError: Метод пока не реализован полностью.
+ """
+ logger.debug(f"Parsing DOC file: {file_path}")
+
+ if not os.path.exists(file_path):
+ raise ValueError(f"File not found: {file_path}")
+
+ filename = os.path.basename(file_path)
+
+ # Создаем заглушку документа
+ doc = ParsedDocument(name=filename, type="DOC")
+
+ # Полная реализация будет добавлена позже
+ # (с использованием textract или pywin32)
+ logger.warning("DOC parsing not fully implemented yet")
+
+ return doc
+
+ def parse(
+ self, file: BinaryIO, file_type: FileType | str | None = None
+ ) -> ParsedDocument:
+ """
+ Парсит DOC документ из объекта файла и возвращает его структурное представление.
+
+ Args:
+ file (BinaryIO): Объект файла для парсинга.
+ file_type: Тип файла, если известен.
+ Может быть объектом FileType или строкой с расширением (".doc").
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ NotImplementedError: Метод пока не реализован полностью.
+ """
+ logger.debug("Parsing DOC from file object")
+
+ if file_type and isinstance(file_type, FileType) and file_type != FileType.DOC:
+ logger.warning(
+ f"Provided file_type {file_type} doesn't match parser type {FileType.DOC}"
+ )
+
+ # Создаем заглушку документа
+ doc = ParsedDocument(name="unknown.doc", type="DOC")
+
+ # Полная реализация будет добавлена позже
+ # (с использованием textract или pywin32)
+ logger.warning("DOC parsing not fully implemented yet")
+
+ return doc
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/__init__.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f778d5c2dc3a3fec81f6997731cfba1981a2299e
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/__init__.py
@@ -0,0 +1,22 @@
+"""
+Подмодуль для работы с DOCX документами.
+
+Содержит компоненты для парсинга различных частей DOCX документов,
+включая стили, метаданные, нумерацию и другие элементы.
+"""
+
+from .core_properties_parser import CorePropertiesParser
+from .metadata_parser import MetadataParser
+from .numbering_parser import NumberingParser
+from .page_estimator import DocxPageEstimator
+from .relationships_parser import RelationshipsParser
+from .styles_parser import StylesParser
+
+__all__ = [
+ "CorePropertiesParser",
+ "MetadataParser",
+ "NumberingParser",
+ "RelationshipsParser",
+ "StylesParser",
+ "DocxPageEstimator",
+]
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/core_properties_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/core_properties_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..328b4dd3ba8566c492d926fb718543db5562dbf9
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/core_properties_parser.py
@@ -0,0 +1,83 @@
+"""
+Модуль для извлечения основных свойств DOCX документа.
+"""
+
+import logging
+import os
+from typing import Any
+
+from bs4 import BeautifulSoup
+
+logger = logging.getLogger(__name__)
+
+
+class CorePropertiesParser:
+ """
+ Парсер для извлечения основных свойств из docProps/core.xml.
+ """
+
+ def parse(self, temp_dir: str) -> dict[str, Any]:
+ """
+ Извлекает основные свойства из docProps/core.xml.
+
+ Args:
+ temp_dir (str): Путь к временной директории с распакованным DOCX.
+
+ Returns:
+ dict[str, Any]: Словарь с основными свойствами.
+ """
+ core_props_path = os.path.join(temp_dir, 'docProps', 'core.xml')
+ if not os.path.exists(core_props_path):
+ logger.warning(f"Core properties file not found: {core_props_path}")
+ return {}
+
+ try:
+ with open(core_props_path, 'rb') as f:
+ content = f.read()
+
+ # Парсим XML с помощью BeautifulSoup
+ soup = BeautifulSoup(content, 'xml')
+
+ # Извлекаем основные свойства
+ props = {}
+
+ # Автор (creator)
+ creator = soup.find('dc:creator')
+ if creator:
+ props['creator'] = creator.text
+
+ # Заголовок (title)
+ title = soup.find('dc:title')
+ if title:
+ props['title'] = title.text
+
+ # Тема (subject)
+ subject = soup.find('dc:subject')
+ if subject:
+ props['subject'] = subject.text
+
+ # Описание (description)
+ description = soup.find('dc:description')
+ if description:
+ props['description'] = description.text
+
+ # Ключевые слова (keywords)
+ keywords = soup.find('cp:keywords')
+ if keywords:
+ props['keywords'] = keywords.text
+
+ # Дата создания (created)
+ created = soup.find('dcterms:created')
+ if created:
+ props['created'] = created.text
+
+ # Дата изменения (modified)
+ modified = soup.find('dcterms:modified')
+ if modified:
+ props['modified'] = modified.text
+
+ return props
+
+ except Exception as e:
+ logger.error(f"Error extracting core properties: {e}")
+ return {}
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/metadata_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/metadata_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..3eb88d85930456e9c3e1ed33b2eabcc81df3c78a
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/metadata_parser.py
@@ -0,0 +1,81 @@
+"""
+Модуль для извлечения метаданных из DOCX документа.
+"""
+
+import logging
+import os
+from typing import Any
+
+from bs4 import BeautifulSoup
+
+logger = logging.getLogger(__name__)
+
+
+class MetadataParser:
+ """
+ Парсер для извлечения метаданных из docProps/app.xml.
+ """
+
+ def parse(self, temp_dir: str) -> dict[str, Any]:
+ """
+ Извлекает метаданные из docProps/app.xml.
+
+ Args:
+ temp_dir (str): Путь к временной директории с распакованным DOCX.
+
+ Returns:
+ dict[str, Any]: Словарь с метаданными.
+ """
+ app_path = os.path.join(temp_dir, 'docProps', 'app.xml')
+ if not os.path.exists(app_path):
+ logger.warning(f"App properties file not found: {app_path}")
+ return {}
+
+ try:
+ with open(app_path, 'rb') as f:
+ content = f.read()
+
+ # Парсим XML с помощью BeautifulSoup
+ soup = BeautifulSoup(content, 'xml')
+
+ # Извлекаем метаданные
+ metadata = {}
+
+ # Статистика документа
+ pages = soup.find('Pages')
+ if pages:
+ metadata['pages'] = int(pages.text)
+
+ words = soup.find('Words')
+ if words:
+ metadata['words'] = int(words.text)
+
+ characters = soup.find('Characters')
+ if characters:
+ metadata['characters'] = int(characters.text)
+
+ # Информация о приложении
+ application = soup.find('Application')
+ if application:
+ metadata['application'] = application.text
+
+ app_version = soup.find('AppVersion')
+ if app_version:
+ metadata['app_version'] = app_version.text
+
+ # Информация о компании
+ company = soup.find('Company')
+ if company:
+ metadata['company'] = company.text
+
+ # Время редактирования
+ total_time = soup.find('TotalTime')
+ if total_time:
+ metadata['total_time'] = int(total_time.text)
+
+ logger.debug("Extracted document metadata")
+ return metadata
+
+ except Exception as e:
+ logger.error(f"Error extracting metadata: {e}")
+ return {}
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/numbering_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/numbering_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b8d2a2516a3f8998630a7b9431052c8564c30f0
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/numbering_parser.py
@@ -0,0 +1,96 @@
+"""
+Модуль для извлечения информации о нумерации из DOCX документа.
+"""
+
+import logging
+import os
+from typing import Any
+
+from bs4 import BeautifulSoup
+
+logger = logging.getLogger(__name__)
+
+
+class NumberingParser:
+ """
+ Парсер для извлечения информации о нумерации из word/numbering.xml.
+ """
+
+ def parse(self, temp_dir: str) -> dict[str, Any]:
+ """
+ Извлекает информацию о нумерации из word/numbering.xml.
+
+ Args:
+ temp_dir (str): Путь к временной директории с распакованным DOCX.
+
+ Returns:
+ dict[str, Any]: Словарь с информацией о нумерации.
+ """
+ numbering_path = os.path.join(temp_dir, 'word', 'numbering.xml')
+ if not os.path.exists(numbering_path):
+ logger.warning(f"Numbering file not found: {numbering_path}")
+ return {}
+
+ try:
+ with open(numbering_path, 'rb') as f:
+ content = f.read()
+
+ # Парсим XML с помощью BeautifulSoup
+ soup = BeautifulSoup(content, 'xml')
+
+ # Извлекаем определения абстрактной нумерации
+ abstract_nums = {}
+ for abstract_num in soup.find_all('w:abstractNum'):
+ if 'w:abstractNumId' in abstract_num.attrs:
+ abstract_id = abstract_num['w:abstractNumId']
+ levels = {}
+
+ # Извлекаем информацию о каждом уровне нумерации
+ for level in abstract_num.find_all('w:lvl'):
+ if 'w:ilvl' in level.attrs:
+ level_id = level['w:ilvl']
+ level_info = {}
+
+ # Формат нумерации (decimal, bullet, etc.)
+ num_fmt = level.find('w:numFmt')
+ if num_fmt and 'w:val' in num_fmt.attrs:
+ level_info['format'] = num_fmt['w:val']
+
+ # Текст до и после номера
+ level_text = level.find('w:lvlText')
+ if level_text and 'w:val' in level_text.attrs:
+ level_info['text'] = level_text['w:val']
+
+ # Выравнивание
+ jc = level.find('w:lvlJc')
+ if jc and 'w:val' in jc.attrs:
+ level_info['alignment'] = jc['w:val']
+
+ levels[level_id] = level_info
+
+ abstract_nums[abstract_id] = levels
+
+ # Извлекаем конкретные определения нумерации
+ numbering = {}
+ for num in soup.find_all('w:num'):
+ if 'w:numId' in num.attrs:
+ num_id = num['w:numId']
+ abstract_num_id = None
+
+ # Получаем ссылку на абстрактную нумерацию
+ abstract_num = num.find('w:abstractNumId')
+ if abstract_num and 'w:val' in abstract_num.attrs:
+ abstract_num_id = abstract_num['w:val']
+
+ if abstract_num_id in abstract_nums:
+ numbering[num_id] = {
+ 'abstract_num_id': abstract_num_id,
+ 'levels': abstract_nums[abstract_num_id]
+ }
+
+ logger.debug(f"Extracted {len(numbering)} numbering definitions")
+ return numbering
+
+ except Exception as e:
+ logger.error(f"Error extracting numbering: {e}")
+ return {}
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/page_estimator.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/page_estimator.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca3f7dc03094ae752affed4d8959dac153e1fb75
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/page_estimator.py
@@ -0,0 +1,188 @@
+"""
+Модуль для оценки номеров страниц в DOCX документах.
+
+Предоставляет приблизительную оценку номера страницы для элементов документа
+на основе метаданных и распределения элементов.
+"""
+
+import logging
+from typing import Dict, Optional, Tuple
+
+from bs4 import BeautifulSoup, Tag
+
+logger = logging.getLogger(__name__)
+
+
+class DocxPageEstimator:
+ """
+ Класс для приблизительной оценки номеров страниц в DOCX документах.
+
+ Использует метаданные документа и равномерное распределение для
+ оценки номеров страниц элементов.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует оценщик страниц DOCX.
+ """
+ pass
+
+ def process_document(self, soup: BeautifulSoup, metadata: Dict = None) -> Tuple[Dict[int, int], Dict[int, int]]:
+ """
+ Быстрый метод оценки страниц документа без сложных вычислений.
+
+ Args:
+ soup (BeautifulSoup): Beautiful Soup объект, содержащий XML документа.
+ metadata (Dict): Метаданные документа, которые могут содержать количество страниц.
+
+ Returns:
+ Tuple[Dict[int, int], Dict[int, int]]: Словари с номерами страниц для параграфов и таблиц.
+ """
+ logger.debug("Using fast page number estimation for DOCX")
+
+ # Получаем общее количество страниц
+ total_pages = 1
+ if metadata and "pages" in metadata:
+ # Используем информацию из метаданных, если она есть
+ total_pages = metadata.get("pages", 1)
+ logger.debug(f"Using page count from metadata: {total_pages}")
+
+ # Находим все параграфы и таблицы
+ paragraphs = soup.find_all("w:p")
+ tables = soup.find_all("w:tbl")
+
+ # Словари для хранения номеров страниц
+ paragraph_pages = {}
+ table_pages = {}
+
+ # Если есть хотя бы одна страница
+ if total_pages > 0:
+ # Распределяем параграфы по страницам равномерно
+ total_paragraphs = len(paragraphs)
+ if total_paragraphs > 0:
+ paragraphs_per_page = max(1, total_paragraphs // total_pages)
+ for i, p in enumerate(paragraphs):
+ page_number = min(total_pages, (i // paragraphs_per_page) + 1)
+ paragraph_pages[i] = page_number
+
+ # Распределяем таблицы по страницам равномерно
+ total_tables = len(tables)
+ if total_tables > 0:
+ tables_per_page = max(1, total_tables // total_pages)
+ for i, t in enumerate(tables):
+ page_number = min(total_pages, (i // tables_per_page) + 1)
+ table_pages[i] = page_number
+
+ return paragraph_pages, table_pages
+
+ def reset(self):
+ """
+ Сбрасывает счетчики страниц.
+ """
+ self.current_page = 1
+ self.chars_on_current_page = 0
+
+ def _process_paragraph(self, paragraph: Tag) -> None:
+ """
+ Обрабатывает параграф и оценивает, на какой странице он находится.
+
+ Args:
+ paragraph (Tag): XML-элемент параграфа.
+ """
+ # Проверяем наличие разрыва страницы в параграфе
+ if paragraph.find('w:br', attrs={'w:type': 'page'}):
+ # Если есть разрыв страницы, переходим на следующую
+ self.current_page += 1
+ self.chars_on_current_page = 0
+ return
+
+ # Получаем текст параграфа
+ text = self._get_paragraph_text(paragraph)
+ text_length = len(text)
+
+ # Проверяем, является ли параграф заголовком
+ style_id = self._get_paragraph_style_id(paragraph)
+ if style_id and ('Heading' in style_id or 'заголовок' in style_id.lower()):
+ # Заголовки занимают больше места
+ text_length *= self.HEADING_SPACE_MULTIPLIER
+
+ # Добавляем длину текста к счетчику символов на текущей странице
+ self.chars_on_current_page += text_length
+
+ # Если превысили лимит символов на странице, переходим на следующую
+ if self.chars_on_current_page > self.DEFAULT_CHARS_PER_PAGE:
+ self.current_page += 1
+ self.chars_on_current_page = self.chars_on_current_page % self.DEFAULT_CHARS_PER_PAGE
+
+ def _process_table(self, table: Tag) -> None:
+ """
+ Обрабатывает таблицу и оценивает, на какой странице она находится.
+
+ Args:
+ table (Tag): XML-элемент таблицы.
+ """
+ # Подсчитываем количество строк в таблице
+ rows = table.find_all('w:tr')
+ total_chars = len(rows) * self.TABLE_ROW_CHARS
+
+ # Добавляем символы таблицы к счетчику
+ self.chars_on_current_page += total_chars
+
+ # Если превысили лимит символов на странице, переходим на следующую
+ while self.chars_on_current_page > self.DEFAULT_CHARS_PER_PAGE:
+ self.current_page += 1
+ self.chars_on_current_page -= self.DEFAULT_CHARS_PER_PAGE
+
+ def _get_paragraph_text(self, paragraph: Tag) -> str:
+ """
+ Извлекает текст из параграфа.
+
+ Args:
+ paragraph (Tag): XML-элемент параграфа.
+
+ Returns:
+ str: Текст параграфа.
+ """
+ # Находим все элементы текста (w:t) внутри параграфа
+ text_elements = paragraph.find_all('w:t')
+ return ''.join(t.get_text() for t in text_elements)
+
+ def _get_paragraph_style_id(self, paragraph: Tag) -> Optional[str]:
+ """
+ Получает идентификатор стиля параграфа.
+
+ Args:
+ paragraph (Tag): XML-элемент параграфа.
+
+ Returns:
+ Optional[str]: Идентификатор стиля или None, если стиль не указан.
+ """
+ # Ищем элемент стиля параграфа
+ para_pr = paragraph.find('w:pPr')
+ if para_pr:
+ style = para_pr.find('w:pStyle')
+ if style and 'w:val' in style.attrs:
+ return style['w:val']
+ return None
+
+ def get_page_number(self, element_id: str, soup: BeautifulSoup) -> Optional[int]:
+ """
+ Возвращает номер страницы для указанного элемента по его ID.
+
+ Args:
+ element_id (str): ID элемента.
+ soup (BeautifulSoup): Beautiful Soup объект, содержащий XML документа.
+
+ Returns:
+ Optional[int]: Номер страницы или None, если элемент не найден.
+ """
+ # Находим элемент по ID
+ element = soup.find(id=element_id)
+ if not element:
+ return None
+
+ # Если элемент имеет атрибут с номером страницы, возвращаем его
+ if 'data-page-number' in element.attrs:
+ return int(element['data-page-number'])
+
+ return None
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/relationships_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/relationships_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..b5a0c8109ef72087c8e9dfe8aea4f024424e6b41
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/relationships_parser.py
@@ -0,0 +1,57 @@
+"""
+Модуль для извлечения связей из DOCX документа.
+"""
+
+import logging
+import os
+from typing import Any
+
+from bs4 import BeautifulSoup
+
+logger = logging.getLogger(__name__)
+
+
+class RelationshipsParser:
+ """
+ Парсер для извлечения связей из word/_rels/document.xml.rels.
+ """
+
+ def parse(self, temp_dir: str) -> dict[str, Any]:
+ """
+ Извлекает связи из word/_rels/document.xml.rels.
+
+ Args:
+ temp_dir (str): Путь к временной директории с распакованным DOCX.
+
+ Returns:
+ dict[str, Any]: Словарь с информацией о связях.
+ """
+ rels_path = os.path.join(temp_dir, 'word', '_rels', 'document.xml.rels')
+ if not os.path.exists(rels_path):
+ logger.warning(f"Relationships file not found: {rels_path}")
+ return {}
+
+ try:
+ with open(rels_path, 'rb') as f:
+ content = f.read()
+
+ # Парсим XML с помощью BeautifulSoup
+ soup = BeautifulSoup(content, 'xml')
+
+ # Извлекаем информацию о связях
+ relationships = {}
+ for relationship in soup.find_all('Relationship'):
+ if all(attr in relationship.attrs for attr in ['Id', 'Type', 'Target']):
+ rel_id = relationship['Id']
+ relationships[rel_id] = {
+ 'type': relationship['Type'].split('/')[-1], # Берем только последнюю часть URI
+ 'target': relationship['Target'],
+ 'target_mode': relationship.get('TargetMode', 'Internal')
+ }
+
+ logger.debug(f"Extracted {len(relationships)} relationships")
+ return relationships
+
+ except Exception as e:
+ logger.error(f"Error extracting relationships: {e}")
+ return {}
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/styles_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/styles_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..7f89435dc57344d74b586ab28cfa388635fa9e48
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx/styles_parser.py
@@ -0,0 +1,74 @@
+"""
+Модуль для извлечения стилей DOCX документа.
+"""
+
+import logging
+import os
+from typing import Any
+
+from bs4 import BeautifulSoup
+
+logger = logging.getLogger(__name__)
+
+
+class StylesParser:
+ """
+ Парсер для извлечения стилей из word/styles.xml.
+ """
+
+ def parse(self, temp_dir: str) -> dict[str, Any]:
+ """
+ Извлекает стили из word/styles.xml.
+
+ Args:
+ temp_dir (str): Путь к временной директории с распакованным DOCX.
+
+ Returns:
+ dict[str, Any]: Словарь с информацией о стилях.
+ """
+ styles_path = os.path.join(temp_dir, 'word', 'styles.xml')
+ if not os.path.exists(styles_path):
+ logger.warning(f"Styles file not found: {styles_path}")
+ return {}
+
+ try:
+ with open(styles_path, 'rb') as f:
+ content = f.read()
+
+ # Парсим XML с помощью BeautifulSoup
+ soup = BeautifulSoup(content, 'xml')
+
+ # Извлекаем информацию о стилях
+ styles_cache = {}
+ for style in soup.find_all('w:style'):
+ if 'w:styleId' in style.attrs:
+ style_id = style['w:styleId']
+ style_info = {}
+
+ # Имя стиля
+ name = style.find('w:name')
+ if name and 'w:val' in name.attrs:
+ style_info['name'] = name['w:val']
+
+ # Тип стиля (paragraph, character, table, numbering)
+ if 'w:type' in style.attrs:
+ style_info['type'] = style['w:type']
+
+ # Базовый стиль
+ base_style = style.find('w:basedOn')
+ if base_style and 'w:val' in base_style.attrs:
+ style_info['based_on'] = base_style['w:val']
+
+ # Следующий стиль
+ next_style = style.find('w:next')
+ if next_style and 'w:val' in next_style.attrs:
+ style_info['next'] = next_style['w:val']
+
+ styles_cache[style_id] = style_info
+
+ logger.debug(f"Extracted {len(styles_cache)} styles")
+ return styles_cache
+
+ except Exception as e:
+ logger.error(f"Error extracting styles: {e}")
+ return {}
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/docx_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd36ce4546692b87f0b0711f0752d10c851669bb
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/docx_parser.py
@@ -0,0 +1,447 @@
+"""
+Модуль для парсинга DOCX документов.
+"""
+
+import io
+import logging
+import os
+import shutil
+import tempfile
+import zipfile
+from typing import Any, BinaryIO
+
+from bs4 import BeautifulSoup
+
+from ...data_classes import ParsedDocument
+from ..abstract_parser import AbstractParser
+from ..file_types import FileType
+from .docx.page_estimator import DocxPageEstimator
+from .docx.relationships_parser import RelationshipsParser
+from .xml_parser import XMLParser
+
+logger = logging.getLogger(__name__)
+
+
+class DocxParser(AbstractParser):
+ """
+ Парсер для DOCX документов.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует парсер DOCX.
+ """
+ super().__init__(FileType.DOCX)
+ self.xml_parser = None
+ self.page_estimator = DocxPageEstimator()
+ self.relationships_parser = RelationshipsParser()
+
+ def parse_by_path(self, file_path: str) -> ParsedDocument:
+ """
+ Парсит DOCX документ по пути к файлу.
+
+ Args:
+ file_path (str): Путь к DOCX файлу.
+
+ Returns:
+ ParsedDocument: Распарсенный документ.
+ """
+ with open(file_path, 'rb') as f:
+ parsed_document = self.parse(f)
+ parsed_document.name = os.path.basename(file_path)
+ parsed_document.type = FileType.DOCX.name
+ return parsed_document
+
+ def parse(self, file: BinaryIO) -> ParsedDocument:
+ """
+ Парсит DOCX документ из файлового объекта.
+
+ Args:
+ file (BinaryIO): Файловый объект DOCX документа.
+
+ Returns:
+ ParsedDocument: Распарсенный документ.
+ """
+ # Создаем временную директорию для распаковки
+ temp_dir = tempfile.mkdtemp()
+ try:
+ # Распаковываем DOCX во временную директорию
+ with zipfile.ZipFile(file) as docx:
+ docx.extractall(temp_dir)
+
+ # Извлекаем стили
+ styles = self._extract_styles(temp_dir)
+
+ # Извлекаем нумерацию
+ numbering = self._extract_numbering(temp_dir)
+
+ # Извлекаем связи
+ relationships = self.relationships_parser.parse(temp_dir)
+
+ # Создаем XML парсер с кэшами
+ self.xml_parser = XMLParser(styles, numbering, relationships)
+
+ # Парсим основное содержимое
+ document_path = os.path.join(temp_dir, 'word', 'document.xml')
+ if not os.path.exists(document_path):
+ logger.error(f"Document file not found: {document_path}")
+ return ParsedDocument([], {})
+
+ # Читаем и парсим основной документ
+ with open(document_path, 'rb') as f:
+ content = f.read()
+
+ # Получаем метаданные
+ metadata = self._extract_metadata(temp_dir)
+
+ # Предварительно оцениваем номера страниц
+ estimated_pages = self._estimate_page_numbers(content)
+
+ # Парсим документ через XMLParser, оборачивая байты в BytesIO
+ doc = self.xml_parser.parse(io.BytesIO(content))
+ doc.meta.note = metadata
+
+ # Применяем номера страниц к элементам документа
+ self._apply_page_numbers(doc, estimated_pages)
+
+ return doc
+
+ finally:
+ # Удаляем временную директорию
+ shutil.rmtree(temp_dir)
+
+ def _extract_styles(self, temp_dir: str) -> dict[str, Any]:
+ """
+ Извлекает стили из word/styles.xml.
+
+ Args:
+ temp_dir (str): Путь к временной директории с распакованным DOCX.
+
+ Returns:
+ dict[str, Any]: Словарь с информацией о стилях.
+ """
+ styles_path = os.path.join(temp_dir, 'word', 'styles.xml')
+ if not os.path.exists(styles_path):
+ logger.warning(f"Styles file not found: {styles_path}")
+ return {}
+
+ try:
+ with open(styles_path, 'rb') as f:
+ content = f.read()
+
+ # Парсим XML с помощью BeautifulSoup
+ soup = BeautifulSoup(content, 'xml')
+
+ # Извлекаем информацию о стилях
+ styles = {}
+ for style in soup.find_all('w:style'):
+ if 'w:styleId' in style.attrs:
+ style_id = style['w:styleId']
+ style_info = {}
+
+ # Имя стиля
+ name = style.find('w:name')
+ if name and 'w:val' in name.attrs:
+ style_info['name'] = name['w:val']
+
+ # Тип стиля (paragraph, character, table, numbering)
+ if 'w:type' in style.attrs:
+ style_info['type'] = style['w:type']
+
+ # Базовый стиль
+ base_style = style.find('w:basedOn')
+ if base_style and 'w:val' in base_style.attrs:
+ style_info['based_on'] = base_style['w:val']
+
+ # Следующий стиль
+ next_style = style.find('w:next')
+ if next_style and 'w:val' in next_style.attrs:
+ style_info['next'] = next_style['w:val']
+
+ styles[style_id] = style_info
+
+ logger.debug(f"Extracted {len(styles)} styles")
+ return styles
+
+ except Exception as e:
+ logger.error(f"Error extracting styles: {e}")
+ return {}
+
+ def _extract_numbering(self, temp_dir: str) -> dict[str, Any]:
+ """
+ Извлекает информацию о нумерации из word/numbering.xml.
+
+ Args:
+ temp_dir (str): Путь к временной директории с распакованным DOCX.
+
+ Returns:
+ dict[str, Any]: Словарь с информацией о нумерации.
+ """
+ numbering_path = os.path.join(temp_dir, 'word', 'numbering.xml')
+ if not os.path.exists(numbering_path):
+ logger.warning(f"Numbering file not found: {numbering_path}")
+ return {}
+
+ try:
+ with open(numbering_path, 'rb') as f:
+ content = f.read()
+
+ # Парсим XML с помощью BeautifulSoup
+ soup = BeautifulSoup(content, 'xml')
+
+ # Извлекаем определения абстрактной нумерации
+ abstract_nums = {}
+ for abstract_num in soup.find_all('w:abstractNum'):
+ if 'w:abstractNumId' in abstract_num.attrs:
+ abstract_id = abstract_num['w:abstractNumId']
+ levels = {}
+
+ # Извлекаем информацию о каждом уровне нумерации
+ for level in abstract_num.find_all('w:lvl'):
+ if 'w:ilvl' in level.attrs:
+ level_id = level['w:ilvl']
+ level_info = {}
+
+ # Формат нумерации (decimal, bullet, etc.)
+ num_fmt = level.find('w:numFmt')
+ if num_fmt and 'w:val' in num_fmt.attrs:
+ level_info['format'] = num_fmt['w:val']
+
+ # Текст до и после номера
+ level_text = level.find('w:lvlText')
+ if level_text and 'w:val' in level_text.attrs:
+ level_info['text'] = level_text['w:val']
+
+ # Выравнивание
+ jc = level.find('w:lvlJc')
+ if jc and 'w:val' in jc.attrs:
+ level_info['alignment'] = jc['w:val']
+
+ levels[level_id] = level_info
+
+ abstract_nums[abstract_id] = levels
+
+ # Извлекаем конкретные определения нумерации
+ numbering = {}
+ for num in soup.find_all('w:num'):
+ if 'w:numId' in num.attrs:
+ num_id = num['w:numId']
+ abstract_num_id = None
+
+ # Получаем ссылку на абстрактную нумерацию
+ abstract_num = num.find('w:abstractNumId')
+ if abstract_num and 'w:val' in abstract_num.attrs:
+ abstract_num_id = abstract_num['w:val']
+
+ if abstract_num_id in abstract_nums:
+ numbering[num_id] = {
+ 'abstract_num_id': abstract_num_id,
+ 'levels': abstract_nums[abstract_num_id],
+ }
+
+ logger.debug(f"Extracted {len(numbering)} numbering definitions")
+ return numbering
+
+ except Exception as e:
+ logger.error(f"Error extracting numbering: {e}")
+ return {}
+
+ def _extract_metadata(self, temp_dir: str) -> dict[str, Any]:
+ """
+ Извлекает метаданные из docProps/core.xml и docProps/app.xml.
+
+ Args:
+ temp_dir (str): Путь к временной директории с распакованным DOCX.
+
+ Returns:
+ dict[str, Any]: Словарь с метаданными.
+ """
+ metadata = {}
+
+ # Извлекаем основные свойства
+ core_props_path = os.path.join(temp_dir, 'docProps', 'core.xml')
+ if os.path.exists(core_props_path):
+ try:
+ with open(core_props_path, 'rb') as f:
+ content = f.read()
+
+ soup = BeautifulSoup(content, 'xml')
+
+ # Автор
+ creator = soup.find('dc:creator')
+ if creator:
+ metadata['creator'] = creator.text
+
+ # Заголовок
+ title = soup.find('dc:title')
+ if title:
+ metadata['title'] = title.text
+
+ # Тема
+ subject = soup.find('dc:subject')
+ if subject:
+ metadata['subject'] = subject.text
+
+ # Описание
+ description = soup.find('dc:description')
+ if description:
+ metadata['description'] = description.text
+
+ # Ключевые слова
+ keywords = soup.find('cp:keywords')
+ if keywords:
+ metadata['keywords'] = keywords.text
+
+ # Даты создания и изменения
+ created = soup.find('dcterms:created')
+ if created:
+ metadata['created'] = created.text
+
+ modified = soup.find('dcterms:modified')
+ if modified:
+ metadata['modified'] = modified.text
+
+ except Exception as e:
+ logger.error(f"Error extracting core properties: {e}")
+
+ # Извлекаем свойства приложения
+ app_props_path = os.path.join(temp_dir, 'docProps', 'app.xml')
+ if os.path.exists(app_props_path):
+ try:
+ with open(app_props_path, 'rb') as f:
+ content = f.read()
+
+ soup = BeautifulSoup(content, 'xml')
+
+ # Статистика документа
+ pages = soup.find('Pages')
+ if pages:
+ metadata['pages'] = int(pages.text)
+
+ words = soup.find('Words')
+ if words:
+ metadata['words'] = int(words.text)
+
+ characters = soup.find('Characters')
+ if characters:
+ metadata['characters'] = int(characters.text)
+
+ # Информация о приложении
+ application = soup.find('Application')
+ if application:
+ metadata['application'] = application.text
+
+ app_version = soup.find('AppVersion')
+ if app_version:
+ metadata['app_version'] = app_version.text
+
+ # Информация о компании
+ company = soup.find('Company')
+ if company:
+ metadata['company'] = company.text
+
+ # Время редактирования
+ total_time = soup.find('TotalTime')
+ if total_time:
+ metadata['total_time'] = int(total_time.text)
+
+ except Exception as e:
+ logger.error(f"Error extracting app properties: {e}")
+
+ # Сохраняем метаданные как атрибут для доступа из других методов
+ self._metadata = metadata
+ return metadata
+
+ def _estimate_page_numbers(self, content: bytes) -> dict[str, int]:
+ """
+ Оценивает номера страниц для элементов документа.
+
+ Args:
+ content (bytes): Содержимое документа.
+
+ Returns:
+ dict[str, int]: Словарь соответствий id элемента и номера страницы.
+ """
+ logger.debug("Estimating page numbers for document elements")
+
+ # Создаем словарь для хранения номеров страниц
+ page_numbers = {}
+
+ try:
+ # Получаем метаданные, включая количество страниц из metadata
+ total_pages = self._metadata.get("pages", 0) if hasattr(self, "_metadata") else 0
+ if total_pages <= 0:
+ total_pages = 1 # Минимум одна страница
+
+ # Парсим XML с помощью BeautifulSoup (это быстрая операция)
+ soup = BeautifulSoup(content, 'xml')
+
+ # Используем упрощенный метод расчета
+ paragraph_pages, table_pages = self.page_estimator.process_document(
+ soup,
+ metadata=self._metadata if hasattr(self, "_metadata") else None
+ )
+
+ # Сохраняем информацию в page_numbers
+ page_numbers['paragraphs'] = paragraph_pages
+ page_numbers['tables'] = table_pages
+ page_numbers['total_pages'] = total_pages
+
+ logger.debug(f"Estimated document has {total_pages} pages")
+ logger.debug(f"Assigned page numbers for {len(paragraph_pages)} paragraphs and {len(table_pages)} tables")
+
+ except Exception as e:
+ logger.error(f"Error estimating page numbers: {e}")
+
+ return page_numbers
+
+ def _apply_page_numbers(self, doc: ParsedDocument, page_numbers: dict[str, int]) -> None:
+ """
+ Применяет оценки номеров страниц к элементам документа.
+
+ Args:
+ doc (ParsedDocument): Документ для обновления.
+ page_numbers (dict[str, int]): Словарь соответствий id элемента и номера страницы.
+ """
+ logger.debug("Applying page numbers to document elements")
+
+ # Получаем информацию о страницах
+ paragraph_pages = page_numbers.get('paragraphs', {})
+ table_pages = page_numbers.get('tables', {})
+ total_pages = page_numbers.get('total_pages', 1)
+
+ logger.debug(f"Applying page numbers: document has {total_pages} pages")
+
+ # Устанавливаем индексы документа и номера страниц для параграфов
+ for i, paragraph in enumerate(doc.paragraphs):
+ # Индекс в документе (хотя это также делается в XMLParser._link_elements)
+ paragraph.index_in_document = i
+
+ # Номер страницы
+ page_num = paragraph_pages.get(i, 1)
+ paragraph.page_number = page_num
+
+ # Устанавливаем индексы и номера страниц для таблиц
+ for i, table in enumerate(doc.tables):
+ # Индекс в документе (хотя это также делается в XMLParser._link_elements)
+ table.index_in_document = i
+
+ # Номер страницы
+ page_num = table_pages.get(i, 1)
+ table.page_number = page_num
+
+ # Для изображений
+ for i, image in enumerate(doc.images):
+ # Индекс в документе (хотя это также делается в XMLParser._link_elements)
+ image.index_in_document = i
+
+ # Номер страницы (примерно)
+ image.page_number = min(total_pages, (i % total_pages) + 1)
+
+ # Для формул
+ for i, formula in enumerate(doc.formulas):
+ # Индекс в документе (хотя это также делается в XMLParser._link_elements)
+ formula.index_in_document = i
+
+ # Номер страницы (примерно)
+ formula.page_number = min(total_pages, (i % total_pages) + 1)
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/email_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/email_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..2368b9059efe9980717152a3800d5c8a09f6c094
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/email_parser.py
@@ -0,0 +1,93 @@
+"""
+Модуль с парсером для почтовых сообщений.
+"""
+
+import logging
+import os
+from typing import BinaryIO
+
+from ...data_classes import ParsedDocument
+from ..abstract_parser import AbstractParser
+from ..file_types import FileType
+
+logger = logging.getLogger(__name__)
+
+
+class EmailParser(AbstractParser):
+ """
+ Парсер для почтовых сообщений (EML).
+
+ Примечание: На данный момент реализация является заглушкой.
+ В будущем будет использоваться библиотека email для EML.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует парсер почтовых сообщений.
+ """
+ super().__init__(FileType.EML)
+
+ def parse_by_path(self, file_path: str) -> ParsedDocument:
+ """
+ Парсит почтовое сообщение по пути к файлу и возвращает его структурное представление.
+
+ Args:
+ file_path (str): Путь к файлу почтового сообщения для парсинга.
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если файл не существует или не может быть прочитан.
+ NotImplementedError: Метод пока не реализован полностью.
+ """
+ logger.debug(f"Parsing email file: {file_path}")
+
+ if not os.path.exists(file_path):
+ raise ValueError(f"File not found: {file_path}")
+
+ filename = os.path.basename(file_path)
+
+ # Создаем заглушку документа
+ doc = ParsedDocument(
+ name=filename,
+ type="EMAIL"
+ )
+
+ # Полная реализация будет добавлена позже
+ # (с использованием библиотеки email для EML)
+ logger.warning("Email parsing not fully implemented yet")
+
+ return doc
+
+ def parse(self, file: BinaryIO, file_type: FileType | str | None = None) -> ParsedDocument:
+ """
+ Парсит почтовое сообщение из объекта файла и возвращает его структурное представление.
+
+ Args:
+ file (BinaryIO): Объект файла для парсинга.
+ file_type: Тип файла, если известен.
+ Может быть объектом FileType или строкой с расширением (".eml").
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ NotImplementedError: Метод пока не реализован полностью.
+ """
+ logger.debug("Parsing email from file object")
+
+ if file_type and isinstance(file_type, FileType) and file_type != FileType.EML:
+ logger.warning(f"Provided file_type {file_type} doesn't match parser type {FileType.EML}")
+
+ # Создаем заглушку документа
+ doc = ParsedDocument(
+ name="unknown.eml",
+ type="EMAIL"
+ )
+
+ # Полная реализация будет добавлена позже
+ # (с использованием библиотеки email для EML)
+ logger.warning("Email parsing not fully implemented yet")
+
+ return doc
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/html_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/html_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..53ea774953f7667b10fd0ceb4a6aaa3eafdbe15a
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/html_parser.py
@@ -0,0 +1,94 @@
+"""
+Модуль с парсером для HTML документов.
+"""
+
+import logging
+import os
+from typing import BinaryIO
+
+
+from ...data_classes import ParsedDocument
+from ..abstract_parser import AbstractParser
+from ..file_types import FileType
+
+logger = logging.getLogger(__name__)
+
+
+class HTMLParser(AbstractParser):
+ """
+ Парсер для HTML документов.
+
+ Примечание: На данный момент реализация является заглушкой.
+ В будущем будет использоваться BeautifulSoup для обработки HTML.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует парсер HTML документов.
+ """
+ super().__init__(FileType.HTML)
+
+ def parse_by_path(self, file_path: str) -> ParsedDocument:
+ """
+ Парсит HTML документ по пути к файлу и возвращает его структурное представление.
+
+ Args:
+ file_path (str): Путь к HTML файлу для парсинга.
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если файл не существует или не может быть прочитан.
+ NotImplementedError: Метод пока не реализован полностью.
+ """
+ logger.debug(f"Parsing HTML file: {file_path}")
+
+ if not os.path.exists(file_path):
+ raise ValueError(f"File not found: {file_path}")
+
+ filename = os.path.basename(file_path)
+
+ # Создаем заглушку документа
+ doc = ParsedDocument(
+ name=filename,
+ type="HTML"
+ )
+
+ # Полная реализация будет добавлена позже
+ # (с использованием BeautifulSoup)
+ logger.warning("HTML parsing not fully implemented yet")
+
+ return doc
+
+ def parse(self, file: BinaryIO, file_type: FileType | str | None = None) -> ParsedDocument:
+ """
+ Парсит HTML документ из объекта файла и возвращает его структурное представление.
+
+ Args:
+ file (BinaryIO): Объект файла для парсинга.
+ file_type: Тип файла, если известен.
+ Может быть объектом FileType или строкой с расширением (".html").
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ NotImplementedError: Метод пока не реализован полностью.
+ """
+ logger.debug("Parsing HTML from file object")
+
+ if file_type and isinstance(file_type, FileType) and file_type != FileType.HTML:
+ logger.warning(f"Provided file_type {file_type} doesn't match parser type {FileType.HTML}")
+
+ # Создаем заглушку документа
+ doc = ParsedDocument(
+ name="unknown.html",
+ type="HTML"
+ )
+
+ # Полная реализация будет добавлена позже
+ # (с использованием BeautifulSoup)
+ logger.warning("HTML parsing not fully implemented yet")
+
+ return doc
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/markdown_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/markdown_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..44bec3a92e82652d7fa754806bafcd86cc2f04ed
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/markdown_parser.py
@@ -0,0 +1,93 @@
+"""
+Модуль с парсером для Markdown документов.
+"""
+
+import logging
+import os
+from typing import BinaryIO
+
+from ...data_classes import ParsedDocument
+from ..abstract_parser import AbstractParser
+from ..file_types import FileType
+
+logger = logging.getLogger(__name__)
+
+
+class MarkdownParser(AbstractParser):
+ """
+ Парсер для Markdown документов.
+
+ Примечание: На данный момент реализация является заглушкой.
+ В будущем будет использоваться библиотека markdown или mistune.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует парсер Markdown документов.
+ """
+ super().__init__(FileType.MD)
+
+ def parse_by_path(self, file_path: str) -> ParsedDocument:
+ """
+ Парсит Markdown документ по пути к файлу и возвращает его структурное представление.
+
+ Args:
+ file_path (str): Путь к Markdown файлу для парсинга.
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если файл не существует или не может быть прочитан.
+ NotImplementedError: Метод пока не реализован полностью.
+ """
+ logger.debug(f"Parsing Markdown file: {file_path}")
+
+ if not os.path.exists(file_path):
+ raise ValueError(f"File not found: {file_path}")
+
+ filename = os.path.basename(file_path)
+
+ # Создаем заглушку документа
+ doc = ParsedDocument(
+ name=filename,
+ type="MARKDOWN"
+ )
+
+ # Полная реализация будет добавлена позже
+ # (с использованием библиотеки markdown или mistune)
+ logger.warning("Markdown parsing not fully implemented yet")
+
+ return doc
+
+ def parse(self, file: BinaryIO, file_type: FileType | str | None = None) -> ParsedDocument:
+ """
+ Парсит Markdown документ из объекта файла и возвращает его структурное представление.
+
+ Args:
+ file (BinaryIO): Объект файла для парсинга.
+ file_type: Тип файла, если известен.
+ Может быть объектом FileType или строкой с расширением (".md").
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ NotImplementedError: Метод пока не реализован полностью.
+ """
+ logger.debug("Parsing Markdown from file object")
+
+ if file_type and isinstance(file_type, FileType) and file_type != FileType.MD:
+ logger.warning(f"Provided file_type {file_type} doesn't match parser type {FileType.MD}")
+
+ # Создаем заглушку документа
+ doc = ParsedDocument(
+ name="unknown.md",
+ type="MARKDOWN"
+ )
+
+ # Полная реализация будет добавлена позже
+ # (с использованием библиотеки markdown или mistune)
+ logger.warning("Markdown parsing not fully implemented yet")
+
+ return doc
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/__init__.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..840c741189c202b24dc890fa5650d919300f5a25
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/__init__.py
@@ -0,0 +1,17 @@
+"""
+Подмодуль для компонентов парсера PDF документов.
+"""
+
+from .formula_parser import PDFFormulaParser
+from .image_parser import PDFImageParser
+from .meta_parser import PDFMetaParser
+from .paragraph_parser import PDFParagraphParser
+from .table_parser import PDFTableParser
+
+__all__ = [
+ "PDFMetaParser",
+ "PDFParagraphParser",
+ "PDFTableParser",
+ "PDFImageParser",
+ "PDFFormulaParser",
+]
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/formula_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/formula_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..549d70f61a60e932200804c273a0ea0c3793e8b9
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/formula_parser.py
@@ -0,0 +1,190 @@
+"""
+Модуль для извлечения формул из PDF-документов.
+"""
+
+import logging
+import re
+
+import fitz # PyMuPDF
+
+from ....data_classes import ParsedFormula
+
+logger = logging.getLogger(__name__)
+
+
+class PDFFormulaParser:
+ """
+ Парсер для извлечения формул из PDF-документов.
+
+ Использует PyMuPDF (fitz) и эвристические методы для поиска и извлечения формул.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует парсер формул PDF.
+ """
+ # Регулярное выражение для поиска потенциальных формул
+ # Ищем строки, содержащие математические символы
+ self.formula_pattern = re.compile(
+ r'(?:[=<>≤≥±×÷√∫∑∏∞≈≠∈∀∃∄∴∵∝∧∨¬→←↔]{1}|'
+ r'\(\s*[\d\w]+\s*[+\-*/]=|\$[^\$]+\$)'
+ )
+
+ # Ищем ссылки на формулы в тексте
+ self.formula_ref_pattern = re.compile(
+ r'(?:формул[аы]|уравнени[еяй])\s*(?:\([^\)]+\)|\d+[\.\d]*)',
+ re.IGNORECASE
+ )
+
+ # Математические символы для эвристики
+ self.math_symbols = set('∫∑∏∞≈≠∈∀∃∄∴∵∝∧∨¬→←↔±×÷√=<>')
+
+ def parse(self, pdf_doc: fitz.Document) -> list[ParsedFormula]:
+ """
+ Извлекает формулы из PDF-документа.
+
+ Args:
+ pdf_doc (fitz.Document): Объект PDF-документа.
+
+ Returns:
+ list[ParsedFormula]: Список извлеченных формул.
+ """
+ logger.debug("Extracting formulas from PDF")
+
+ result = []
+
+ # Обходим все страницы документа
+ for page_idx in range(len(pdf_doc)):
+ page = pdf_doc[page_idx]
+ page_num = page_idx + 1
+ logger.debug(f"Processing page {page_num} for formulas")
+
+ try:
+ # Извлекаем блоки текста со страницы
+ blocks = page.get_text("blocks")
+
+ # Ищем потенциальные формулы в каждом блоке
+ for block_idx, block in enumerate(blocks):
+ # block[4] содержит текст блока
+ text = block[4].strip()
+
+ if not self._is_potential_formula(text):
+ continue
+
+ # Ищем формулу в блоке текста
+ formula_match = self.formula_pattern.search(text)
+ if formula_match:
+ # Ищем заголовок формулы
+ formula_title = self._find_formula_title(page, block)
+
+ # Создаем объект формулы
+ formula = ParsedFormula(
+ title=formula_title,
+ latex=text,
+ formula_number=self._extract_formula_number(text),
+ page_number=page_num,
+ index_in_document=len(result)
+ )
+
+ result.append(formula)
+ logger.debug(f"Found potential formula: {text}")
+
+ except Exception as e:
+ logger.error(f"Error analyzing page {page_num} for formulas: {e}")
+ logger.exception(e)
+
+ logger.debug(f"Extracted {len(result)} potential formulas from PDF")
+ return result
+
+ def _is_potential_formula(self, text: str) -> bool:
+ """
+ Проверяет, является ли текст потенциальной формулой.
+
+ Args:
+ text (str): Текст для проверки.
+
+ Returns:
+ bool: True, если текст может быть формулой.
+ """
+ # Пустой текст не является формулой
+ if not text:
+ return False
+
+ # Слишком длинный текст, вероятно, не является формулой
+ if len(text) > 300:
+ return False
+
+ # Содержит математические символы
+ if any(char in self.math_symbols for char in text):
+ return True
+
+ # Содержит равенства или выражения типа "a + b = c"
+ if '=' in text and any(op in text for op in ['+', '-', '*', '/', '^']):
+ return True
+
+ # LaTeX-подобные формулы в долларах
+ if '$' in text:
+ return True
+
+ return False
+
+ def _find_formula_title(self, page: fitz.Page, block: tuple) -> str | None:
+ """
+ Ищет заголовок формулы в тексте вокруг блока.
+
+ Args:
+ page (fitz.Page): Страница PDF-документа.
+ block (tuple): Блок текста с формулой.
+
+ Returns:
+ str | None: Заголовок формулы или None.
+ """
+ # Получаем координаты блока
+ x0, y0, x1, y1 = block[:4]
+
+ # Создаем прямоугольник вокруг блока, расширенный для поиска контекста
+ context_rect = fitz.Rect(
+ max(0, x0 - 50), # Расширение влево
+ max(0, y0 - 50), # Расширение вверх
+ min(page.rect.width, x1 + 50), # Расширение вправо
+ min(page.rect.height, y1 + 50) # Расширение вниз
+ )
+
+ # Извлекаем текст из контекстного прямоугольника
+ context_text = page.get_text("text", clip=context_rect)
+
+ # Ищем упоминание формулы в тексте
+ match = self.formula_ref_pattern.search(context_text)
+ if match:
+ # Пытаемся извлечь полную строку с упоминанием формулы
+ lines = context_text.split('\n')
+ for line in lines:
+ if match.group(0) in line:
+ return line.strip()
+
+ # Если не нашли полную строку, возвращаем само совпадение
+ return match.group(0).strip()
+
+ return None
+
+ def _extract_formula_number(self, text: str) -> str | None:
+ """
+ Извлекает номер формулы из текста.
+
+ Args:
+ text (str): Текст формулы.
+
+ Returns:
+ str | None: Номер формулы или None.
+ """
+ # Ищем номер формулы в скобках (часто формат "(1)" или "(1.2)")
+ bracket_match = re.search(r'\((\d+[\.\d]*)\)', text)
+ if bracket_match:
+ return bracket_match.group(1)
+
+ # Ищем просто числа, если они выделены в тексте
+ number_match = re.search(r'(? list[ParsedImage]:
+ """
+ Извлекает изображения из PDF-документа.
+
+ Args:
+ pdf_doc (fitz.Document): Объект PDF-документа.
+
+ Returns:
+ list[ParsedImage]: Список извлеченных изображений.
+ """
+ logger.debug("Extracting images from PDF")
+
+ result = []
+ image_count = 0
+
+ # Обходим все страницы документа
+ for page_idx in range(len(pdf_doc)):
+ page = pdf_doc[page_idx]
+ page_num = page_idx + 1
+ logger.debug(f"Processing page {page_num} for images")
+
+ try:
+ # PyMuPDF позволяет извлечь изображения со страницы
+ image_list = page.get_images(full=True)
+
+ # Обрабатываем каждое изображение
+ for img_idx, img_info in enumerate(image_list):
+ image_count += 1
+ try:
+ # Извлекаем и создаем объект изображения
+ image = self._extract_image(pdf_doc, img_info, page_num, len(result))
+ if image:
+ result.append(image)
+ except Exception as e:
+ img_id = img_info[0] if len(img_info) > 0 else "unknown"
+ logger.error(f"Error extracting image {img_id} from page {page_num}: {e}")
+ logger.exception(e)
+
+ except Exception as e:
+ logger.error(f"Error processing page {page_num} for images: {e}")
+ logger.exception(e)
+
+ logger.debug(f"Found {image_count} images, successfully extracted {len(result)} images")
+ return result
+
+ def _extract_image(self, pdf_doc: fitz.Document, img_info: tuple, page_num: int, index: int) -> ParsedImage | None:
+ """
+ Извлекает изображение из документа PDF.
+
+ Args:
+ pdf_doc (fitz.Document): Документ PDF.
+ img_info (tuple): Информация об изображении.
+ page_num (int): Номер страницы.
+ index (int): Индекс изображения в документе.
+
+ Returns:
+ ParsedImage | None: Извлеченное изображение или None в случае ошибки.
+ """
+ try:
+ # Получаем основную информацию об изображении
+ img_id = img_info[0] # xref (уникальный идентификатор) изображения
+ img_name = f"Image_{page_num}_{img_id}"
+
+ # Получаем базовое изображение
+ base_image = pdf_doc.extract_image(img_id)
+
+ if base_image:
+ # Получаем данные изображения, размеры и формат
+ image_data = base_image["image"]
+ width = base_image.get("width", 0)
+ height = base_image.get("height", 0)
+
+ # Создаем объект изображения
+ image = ParsedImage(
+ title=img_name,
+ image=image_data,
+ width=width,
+ height=height,
+ page_number=page_num,
+ index_in_document=index
+ )
+ return image
+ else:
+ logger.warning(f"Failed to extract image data for {img_name}")
+ return None
+
+ except Exception as e:
+ logger.error(f"Error extracting image: {e}")
+ logger.exception(e)
+ return None
+
+ def _find_image_caption(self, page: fitz.Page, img_bbox: tuple) -> str | None:
+ """
+ Пытается найти подпись к изображению.
+
+ Args:
+ page (fitz.Page): Страница PDF.
+ img_bbox (tuple): Координаты границ изображения.
+
+ Returns:
+ str | None: Найденная подпись или None.
+ """
+ # Примечание: эта функция пока не используется, но может быть полезна
+ # для будущего улучшения обработки изображений
+ try:
+ # Создаем прямоугольник под изображением, где может быть подпись
+ x0, y0, x1, y1 = img_bbox
+ caption_rect = fitz.Rect(x0, y1, x1, y1 + 50) # 50 единиц ниже изображения
+
+ # Получаем текст в этом прямоугольнике
+ caption_text = page.get_text("text", clip=caption_rect)
+
+ # Проверяем, содержит ли текст ключевые слова, характерные для подписей
+ if caption_text and any(keyword in caption_text.lower()
+ for keyword in ["рис", "рисунок", "изображение", "figure", "pic"]):
+ return caption_text.strip()
+
+ return None
+ except Exception as e:
+ logger.warning(f"Error finding image caption: {e}")
+ return None
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/meta_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/meta_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..9c78606f759930fbc9bbbec83ccfbb0d2a1d3d5e
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/meta_parser.py
@@ -0,0 +1,168 @@
+"""
+Модуль для извлечения метаданных из PDF-документов.
+"""
+
+import logging
+import os
+from typing import Any
+
+import fitz # PyMuPDF
+
+logger = logging.getLogger(__name__)
+
+
+class PDFMetaParser:
+ """
+ Парсер для извлечения метаданных из PDF-документов.
+
+ Использует PyMuPDF (fitz) для доступа к метаданным PDF.
+ """
+
+ def parse(self, pdf_doc: fitz.Document, filepath: str | None = None) -> dict[str, Any]:
+ """
+ Извлекает метаданные из PDF-документа.
+
+ Args:
+ pdf_doc (fitz.Document): Объект PDF-документа.
+ filepath (str | None): Путь к файлу, если доступен.
+
+ Returns:
+ dict[str, Any]: Словарь с метаданными документа.
+ """
+ result = {}
+
+ try:
+ # Базовая информация о файле
+ if filepath:
+ result["filename"] = os.path.basename(filepath)
+ result["filepath"] = filepath
+ # Получаем размер файла
+ result["filesize"] = os.path.getsize(filepath)
+
+ # Извлекаем метаданные документа
+ metadata = pdf_doc.metadata
+
+ if metadata:
+ # Основные метаданные PDF
+ for key in ["title", "author", "subject", "keywords", "creator", "producer", "creationDate", "modDate"]:
+ if key in metadata and metadata[key]:
+ # Преобразуем ключи в более человекочитаемый формат
+ readable_key = key
+ if key == "creationDate":
+ readable_key = "creation_date"
+ elif key == "modDate":
+ readable_key = "modification_date"
+
+ result[readable_key] = metadata[key]
+
+ # Добавляем информацию о структуре документа
+ result["page_count"] = len(pdf_doc)
+
+ # Проверяем наличие закладок (оглавления)
+ toc = pdf_doc.get_toc()
+ if toc:
+ result["has_toc"] = True
+ result["toc_items"] = len(toc)
+ else:
+ result["has_toc"] = False
+
+ # PDF версия
+ if hasattr(pdf_doc, "pdf_version"):
+ result["pdf_version"] = pdf_doc.pdf_version
+
+ # Проверка защиты документа в новой версии PyMuPDF
+ result["is_encrypted"] = pdf_doc.is_encrypted
+ if pdf_doc.is_encrypted:
+ # Извлекаем права доступа с использованием нового API
+ permissions = self._get_permissions(pdf_doc)
+ if permissions:
+ result["permissions"] = permissions
+
+ # Получаем размер страницы (формат) первой страницы
+ if len(pdf_doc) > 0:
+ first_page = pdf_doc[0]
+ page_size = first_page.rect
+ result["page_width"] = page_size.width
+ result["page_height"] = page_size.height
+
+ # Определяем формат страницы (А4, Letter и т.д.)
+ result["page_format"] = self._detect_page_format(page_size.width, page_size.height)
+
+ except Exception as e:
+ logger.error(f"Error extracting PDF metadata: {e}")
+ logger.exception(e)
+ result["error"] = str(e)
+
+ return result
+
+ def _get_permissions(self, pdf_doc: fitz.Document) -> dict[str, bool]:
+ """
+ Получает информацию о правах доступа к документу.
+
+ Args:
+ pdf_doc (fitz.Document): Объект PDF-документа.
+
+ Returns:
+ dict[str, bool]: Словарь с правами доступа.
+ """
+ permissions = {}
+
+ # Используем методы PyMuPDF для проверки прав доступа
+ try:
+ # В новой версии PyMuPDF права доступа доступны через permission_flags
+ if hasattr(pdf_doc, "permission_flags"):
+ perm_flags = pdf_doc.permission_flags
+ permissions["can_print"] = bool(perm_flags & fitz.PDF_PERM_PRINT)
+ permissions["can_copy"] = bool(perm_flags & fitz.PDF_PERM_COPY)
+ permissions["can_modify"] = bool(perm_flags & fitz.PDF_PERM_MODIFY)
+ permissions["can_annotate"] = bool(perm_flags & fitz.PDF_PERM_ANNOTATE)
+ permissions["can_fill_forms"] = bool(perm_flags & fitz.PDF_PERM_FORM)
+ else:
+ # Альтернативный метод в случае отсутствия permission_flags
+ permissions["can_print"] = True
+ permissions["can_copy"] = True
+ permissions["can_modify"] = True
+ permissions["can_annotate"] = True
+ permissions["can_fill_forms"] = True
+ except Exception as e:
+ logger.warning(f"Error getting permissions: {e}")
+
+ return permissions
+
+ def _detect_page_format(self, width: float, height: float) -> str:
+ """
+ Определяет формат страницы на основе её размеров.
+
+ Args:
+ width (float): Ширина страницы в точках.
+ height (float): Высота страницы в точках.
+
+ Returns:
+ str: Название формата страницы или "Custom".
+ """
+ # Преобразуем точки (pt) в миллиметры (мм)
+ mm_width = width * 0.352778
+ mm_height = height * 0.352778
+
+ # Стандартные форматы бумаги (в мм)
+ formats = {
+ "A4": (210, 297),
+ "A3": (297, 420),
+ "A5": (148, 210),
+ "Letter": (215.9, 279.4),
+ "Legal": (215.9, 355.6)
+ }
+
+ # Проверяем известные форматы с допустимым отклонением
+ tolerance = 3 # допустимое отклонение в мм
+
+ for format_name, (format_width, format_height) in formats.items():
+ # Проверяем как в портретной, так и в альбомной ориентации
+ if (abs(mm_width - format_width) <= tolerance and abs(mm_height - format_height) <= tolerance) or \
+ (abs(mm_width - format_height) <= tolerance and abs(mm_height - format_width) <= tolerance):
+ # Определяем ориентацию
+ orientation = "portrait" if mm_height >= mm_width else "landscape"
+ return f"{format_name} ({orientation})"
+
+ # Если не найдено совпадений со стандартными форматами
+ return f"Custom ({mm_width:.1f} x {mm_height:.1f} mm)"
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/paragraph_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/paragraph_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..4cbac6328abfa320989f80a04ec88df054fee053
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/paragraph_parser.py
@@ -0,0 +1,193 @@
+"""
+Модуль для извлечения текстовых блоков из PDF-документов.
+"""
+
+import logging
+import re
+
+import fitz # PyMuPDF
+
+from ....data_classes import ParsedTextBlock, TextStyle
+
+logger = logging.getLogger(__name__)
+
+
+class PDFParagraphParser:
+ """
+ Парсер для извлечения текстовых блоков из PDF-документов.
+
+ Использует PyMuPDF (fitz) для извлечения текста из страниц PDF.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует парсер параграфов PDF.
+ """
+ self.paragraph_pattern = re.compile(r'\n\s*\n')
+
+ def parse(self, pdf_doc: fitz.Document) -> list[ParsedTextBlock]:
+ """
+ Извлекает текстовые блоки из PDF-документа.
+
+ Args:
+ pdf_doc (fitz.Document): Объект PDF-документа.
+
+ Returns:
+ list[ParsedTextBlock]: Список извлеченных текстовых блоков.
+ """
+ logger.debug("Extracting paragraphs from PDF")
+
+ result = []
+
+ # Обходим все страницы документа
+ for page_idx in range(len(pdf_doc)):
+ page = pdf_doc[page_idx]
+ page_num = page_idx + 1
+ logger.debug(f"Processing page {page_num}")
+
+ try:
+ # PyMuPDF позволяет извлекать блоки текста с сохранением структуры
+ blocks = page.get_text("blocks")
+
+ # Обрабатываем каждый блок текста как потенциальный параграф
+ for block_idx, block in enumerate(blocks):
+ # block[4] содержит текст блока
+ text = block[4].strip()
+
+ if not text:
+ continue
+
+ # Разбиваем текст на параграфы
+ paragraphs = self._split_to_paragraphs(text)
+
+ for j, paragraph_text in enumerate(paragraphs):
+ if not paragraph_text.strip():
+ continue
+
+ # Создаем текстовый блок
+ paragraph = ParsedTextBlock(
+ text=paragraph_text.strip(),
+ style=self._detect_style(paragraph_text, page, block),
+ metadata=[{
+ "page": page_num,
+ "block_index": block_idx,
+ "paragraph_index": j,
+ "text": paragraph_text.strip(),
+ "bbox": block[:4] # координаты блока x0, y0, x1, y1
+ }],
+ page_number=page_num,
+ index_in_document=len(result)
+ )
+
+ result.append(paragraph)
+
+ except Exception as e:
+ logger.error(f"Error extracting text from page {page_num}: {e}")
+ logger.exception(e)
+
+ logger.debug(f"Extracted {len(result)} paragraphs from PDF")
+ return result
+
+ def _split_to_paragraphs(self, text: str) -> list[str]:
+ """
+ Разбивает текст на параграфы.
+
+ Args:
+ text (str): Текст для разбиения.
+
+ Returns:
+ list[str]: Список параграфов.
+ """
+ # Разбиваем по пустым строкам, сохраняя пробельные символы
+ return self.paragraph_pattern.split(text)
+
+ def _detect_style(self, text: str, page: fitz.Page, block: tuple) -> TextStyle:
+ """
+ Определяет стиль текста на основе его содержимого и форматирования.
+
+ Args:
+ text (str): Текст параграфа.
+ page (fitz.Page): Объект страницы PyMuPDF.
+ block (tuple): Блок текста из PyMuPDF.
+
+ Returns:
+ TextStyle: Определенный стиль текста.
+ """
+ style = TextStyle()
+
+ try:
+ # Координаты блока текста
+ x0, y0, x1, y1 = block[:4]
+
+ # Определяем, является ли текст заголовком по размеру шрифта
+ words = page.get_text("words", clip=(x0, y0, x1, y1))
+
+ if words:
+ # Проверяем шрифт и стиль
+ fonts = []
+ sizes = []
+ is_bold = False
+ is_italic = False
+
+ # В новых версиях PyMuPDF метод get_texttrace не принимает аргументов
+ # Получаем всю информацию о шрифтах на странице
+ try:
+ font_info = page.get_texttrace()
+
+ # Фильтруем информацию о шрифтах для нашего блока
+ for span in font_info:
+ # Проверяем, находится ли спан в границах нашего блока
+ span_rect = fitz.Rect(span.get("bbox", [0, 0, 0, 0]))
+ block_rect = fitz.Rect(x0, y0, x1, y1)
+
+ if span_rect.intersects(block_rect) and span.get("font"):
+ font_name = span["font"].lower()
+ fonts.append(font_name)
+ sizes.append(span.get("size", 0))
+
+ # Проверяем, содержит ли имя шрифта "bold" или "italic"
+ if "bold" in font_name:
+ is_bold = True
+ if "italic" in font_name or "oblique" in font_name:
+ is_italic = True
+ except Exception as font_err:
+ logger.debug(f"Cannot get detailed font info: {font_err}")
+ # Если не удалось получить информацию о шрифтах, пытаемся угадать
+ # стиль на основе самого текста
+ is_bold = text.isupper() or any(line.strip().startswith('#') for line in text.split('\n'))
+
+ # Определяем, является ли текст заголовком
+ if fonts and sizes:
+ # Если размер шрифта больше среднего или текст полностью в верхнем регистре
+ avg_size = sum(sizes) / len(sizes) if sizes else 0
+ if avg_size > 12 or text.upper() == text:
+ style.paragraph_style = "heading"
+
+ # Устанавливаем флаги стиля текста
+ style.fully_bold = is_bold
+ style.fully_italic = is_italic
+
+ # Определяем, является ли текст списком
+ lines = text.split('\n')
+ if any(line.strip().startswith(('•', '-', '*', '✓', '✔', '✗', '✘', '1.', '2.', 'a.', 'A.'))
+ for line in lines):
+ style.has_numbering = True
+
+ # Определяем выравнивание на основе позиции текста на странице
+ page_width = page.rect.width
+ center = page_width / 2
+
+ # Примерное определение выравнивания
+ if abs(x0 - page.rect.x0) < 50 and abs(x1 - page.rect.x1) < 50:
+ style.alignment = "justified"
+ elif abs(x0 - page.rect.x0) < 50:
+ style.alignment = "left"
+ elif abs(x1 - page.rect.x1) < 50:
+ style.alignment = "right"
+ elif abs((x0 + x1) / 2 - center) < 50:
+ style.alignment = "center"
+
+ except Exception as e:
+ logger.warning(f"Error detecting style: {e}")
+
+ return style
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/table_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/table_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb3c2d48ca552bace8ca0e44f9c36e554318a11c
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf/table_parser.py
@@ -0,0 +1,304 @@
+"""
+Модуль для извлечения таблиц из PDF-документов.
+"""
+
+import logging
+import re
+from typing import Any
+
+import fitz # PyMuPDF
+
+from ....data_classes import ParsedRow, ParsedSubtable, ParsedTable, TextStyle
+
+logger = logging.getLogger(__name__)
+
+
+class PDFTableParser:
+ """
+ Парсер для извлечения таблиц из PDF-документов.
+
+ Использует PyMuPDF (fitz) для анализа структуры документа и определения таблиц.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует парсер таблиц PDF.
+ """
+ # Регулярное выражение для поиска потенциальных заголовков таблиц
+ self.table_header_pattern = re.compile(
+ r'(таблица|табл\.)\s*(\d+|[IVX]+)(\.[\d]+)?', re.IGNORECASE
+ )
+
+ def parse(self, pdf_doc: fitz.Document) -> list[ParsedTable]:
+ """
+ Извлекает таблицы из PDF-документа.
+
+ Args:
+ pdf_doc (fitz.Document): Объект PDF-документа.
+
+ Returns:
+ list[ParsedTable]: Список извлеченных таблиц.
+ """
+ logger.debug("Extracting tables from PDF")
+
+ result = []
+
+ # Обходим все страницы документа
+ for page_idx in range(len(pdf_doc)):
+ page = pdf_doc[page_idx]
+ page_num = page_idx + 1
+ logger.debug(f"Processing page {page_num} for tables")
+
+ try:
+ # В PyMuPDF 1.21+ find_tables() возвращает объект TableFinder, у которого есть свойство tables
+ table_finder = page.find_tables()
+
+ # Проверяем наличие таблиц в TableFinder
+ if hasattr(table_finder, 'tables') and table_finder.tables:
+ for idx, table in enumerate(table_finder.tables):
+ try:
+ # Извлекаем заголовок таблицы из текста над ней
+ title = self._find_table_title(page, table, idx)
+
+ # Извлекаем данные таблицы
+ rows_data = table.extract()
+
+ # Создаем таблицу только если есть данные
+ if rows_data and len(rows_data) > 0:
+ # Подготавливаем данные для создания ParsedTable
+ table_info = {
+ "title": title if title else f"Таблица {page_num}.{idx+1}",
+ "page": page_num,
+ "rows": rows_data,
+ "bbox": table.bbox if hasattr(table, 'bbox') else None
+ }
+
+ # Создаем объект таблицы
+ parsed_table = ParsedTable(
+ title=table_info["title"],
+ subtables=[self._create_subtable(table_info)],
+ page_number=page_num,
+ index_in_document=len(result)
+ )
+ result.append(parsed_table)
+ except Exception as e:
+ logger.error(f"Error creating table object on page {page_num}: {e}")
+ logger.exception(e)
+
+ # Если табличный анализатор не нашел таблиц, попробуем использовать эвристический метод
+ if not hasattr(table_finder, 'tables') or not table_finder.tables:
+ # Получаем текст страницы
+ page_text = page.get_text("text")
+ tables_info = self._detect_tables_heuristically(page_text, page_num)
+
+ for table_info in tables_info:
+ try:
+ # Создаем объект таблицы
+ table = ParsedTable(
+ title=table_info["title"],
+ subtables=[self._create_subtable(table_info)],
+ page_number=page_num,
+ index_in_document=len(result)
+ )
+ result.append(table)
+ except Exception as e:
+ logger.error(f"Error creating table object from heuristic detection on page {page_num}: {e}")
+ logger.exception(e)
+
+ except Exception as e:
+ logger.error(f"Error processing page {page_num} for tables: {e}")
+ logger.exception(e)
+
+ logger.debug(f"Extracted {len(result)} tables from PDF")
+ return result
+
+ def _find_table_title(self, page: fitz.Page, table: Any, table_idx: int) -> str | None:
+ """
+ Ищет заголовок таблицы в тексте над ней.
+
+ Args:
+ page (fitz.Page): Страница PDF-документа.
+ table (Any): Объект таблицы (из find_tables).
+ table_idx (int): Индекс таблицы на странице.
+
+ Returns:
+ str | None: Найденный заголовок таблицы или None.
+ """
+ # Получаем прямоугольник над таблицей
+ # Обрабатываем случай, когда bbox может быть и объектом Rect, и кортежем
+ if hasattr(table, 'bbox'):
+ # Это объект с атрибутом bbox
+ bbox = table.bbox
+ if isinstance(bbox, tuple) or isinstance(bbox, list):
+ # bbox - это кортеж или список координат [x0, y0, x1, y1]
+ x0, y0, x1, y1 = bbox
+ else:
+ # bbox - это объект Rect
+ x0, y0, x1, y1 = bbox.x0, bbox.y0, bbox.x1, bbox.y1
+ elif hasattr(table, 'rect'):
+ # Это объект с атрибутом rect
+ rect = table.rect
+ if isinstance(rect, tuple) or isinstance(rect, list):
+ # rect - это кортеж или список координат [x0, y0, x1, y1]
+ x0, y0, x1, y1 = rect
+ else:
+ # rect - это объект Rect
+ x0, y0, x1, y1 = rect.x0, rect.y0, rect.x1, rect.y1
+ else:
+ # Если мы не можем получить границы таблицы, используем значения по умолчанию
+ # для всей страницы
+ x0, y0, x1, y1 = page.rect.x0, page.rect.y0, page.rect.x1, page.rect.y1
+
+ # Создаем прямоугольник над таблицей
+ above_rect = fitz.Rect(
+ x0,
+ max(0, y0 - 100), # Проверяем до 100 единиц выше таблицы
+ x1,
+ y0
+ )
+
+ # Извлекаем текст над таблицей
+ above_text = page.get_text("text", clip=above_rect)
+
+ # Ищем заголовок таблицы с помощью регулярного выражения
+ match = self.table_header_pattern.search(above_text)
+ if match:
+ # Ищем полную строку с заголовком
+ lines = above_text.split('\n')
+ for line in lines:
+ if match.group(0) in line:
+ return line.strip()
+
+ # Если не нашли полную строку, возвращаем весь текст над таблицей
+ return above_text.strip()
+
+ # Если не нашли явного заголовка, проверяем по ключевым словам
+ if "таблица" in above_text.lower() or "табл." in above_text.lower():
+ # Возвращаем первую строку с упоминанием таблицы
+ lines = above_text.split('\n')
+ for line in lines:
+ if "таблица" in line.lower() or "табл." in line.lower():
+ return line.strip()
+
+ # Если не нашли заголовка, возвращаем None
+ return None
+
+ def _detect_tables_heuristically(self, page_text: str, page_num: int) -> list[dict[str, Any]]:
+ """
+ Использует эвристический метод для определения таблиц в тексте страницы.
+
+ Args:
+ page_text (str): Текст страницы.
+ page_num (int): Номер страницы.
+
+ Returns:
+ list[dict[str, Any]]: Список информации о найденных таблицах.
+ """
+ tables_info = []
+ lines = page_text.split('\n')
+ current_table = None
+
+ for line_idx, line in enumerate(lines):
+ # Проверяем, содержит ли строка заголовок таблицы
+ table_header_match = self.table_header_pattern.search(line)
+
+ if table_header_match:
+ # Если нашли новый заголовок таблицы, сохраняем предыдущую таблицу
+ if current_table and current_table["rows"]:
+ tables_info.append(current_table)
+
+ # Начинаем новую таблицу
+ current_table = {
+ "title": line.strip(),
+ "page": page_num,
+ "rows": []
+ }
+
+ # Ищем строки до следующего заголовка или до конца страницы
+ for next_line_idx in range(line_idx + 1, len(lines)):
+ next_line = lines[next_line_idx].strip()
+
+ # Прекращаем, если нашли новый заголовок таблицы
+ if self.table_header_pattern.search(next_line):
+ break
+
+ # Пропускаем пустые строки
+ if not next_line:
+ continue
+
+ # Эвристика: строки таблицы часто содержат много пробелов или табуляций
+ if ' ' in next_line or '\t' in next_line or len(next_line.split()) >= 3:
+ cells = self._split_line_to_cells(next_line)
+ if cells:
+ current_table["rows"].append(cells)
+
+ # Добавляем последнюю таблицу, если она есть
+ if current_table and current_table["rows"]:
+ tables_info.append(current_table)
+
+ return tables_info
+
+ def _split_line_to_cells(self, line: str) -> list[str]:
+ """
+ Разбивает строку на ячейки таблицы.
+
+ Args:
+ line (str): Строка для разбиения.
+
+ Returns:
+ list[str]: Список ячеек.
+ """
+ # Пробуем разделить по нескольким пробелам или табуляции
+ if '\t' in line:
+ return [cell.strip() for cell in line.split('\t') if cell.strip()]
+
+ # Если есть последовательности пробелов, разбиваем по ним
+ splits = re.split(r'\s{2,}', line)
+ if len(splits) > 1:
+ return [cell.strip() for cell in splits if cell.strip()]
+
+ # Иначе просто разбиваем по одиночным пробелам
+ return [cell.strip() for cell in line.split() if cell.strip()]
+
+ def _create_subtable(self, table_info: dict[str, Any]) -> ParsedSubtable:
+ """
+ Создает объект подтаблицы из информации о таблице.
+
+ Args:
+ table_info (dict[str, Any]): Информация о таблице.
+
+ Returns:
+ ParsedSubtable: Созданная подтаблица.
+ """
+ rows = []
+
+ # Первую строку считаем заголовком, если она есть
+ header = None
+ if table_info["rows"]:
+ header_cells = table_info["rows"][0]
+ header = ParsedRow(
+ cells=header_cells,
+ is_header=True,
+ style=TextStyle(fully_bold=True)
+ )
+
+ # Создаем остальные строки
+ for i, row_cells in enumerate(table_info["rows"][1:], start=1):
+ # Если количество ячеек не совпадает с заголовком, корректируем
+ if len(row_cells) < len(header_cells):
+ row_cells.extend([""] * (len(header_cells) - len(row_cells)))
+ elif len(row_cells) > len(header_cells):
+ row_cells = row_cells[:len(header_cells)]
+
+ row = ParsedRow(
+ index=i,
+ cells=row_cells
+ )
+ rows.append(row)
+
+ # Создаем подтаблицу
+ return ParsedSubtable(
+ title=table_info["title"],
+ header=header,
+ rows=rows
+ )
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..963fb9a28c586dd3f1a51bd67c5b33ab42256375
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/pdf_parser.py
@@ -0,0 +1,228 @@
+"""
+Модуль с парсером для PDF документов.
+"""
+
+import io
+import logging
+import os
+from typing import BinaryIO
+
+import fitz # PyMuPDF
+
+from ...data_classes import ParsedDocument, ParsedMeta
+from ..abstract_parser import AbstractParser
+from ..file_types import FileType
+from .pdf.formula_parser import PDFFormulaParser
+from .pdf.image_parser import PDFImageParser
+from .pdf.meta_parser import PDFMetaParser
+from .pdf.paragraph_parser import PDFParagraphParser
+from .pdf.table_parser import PDFTableParser
+
+logger = logging.getLogger(__name__)
+
+
+class PDFParser(AbstractParser):
+ """
+ Парсер для PDF документов.
+
+ Использует PyMuPDF (fitz) для извлечения текста, изображений, таблиц,
+ формул и метаданных из документа.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует PDF парсер и его компоненты.
+ """
+ super().__init__(FileType.PDF)
+ self.meta_parser = PDFMetaParser()
+ self.paragraph_parser = PDFParagraphParser()
+ self.table_parser = PDFTableParser()
+ self.image_parser = PDFImageParser()
+ self.formula_parser = PDFFormulaParser()
+
+ def parse_by_path(self, file_path: str) -> ParsedDocument:
+ """
+ Парсит PDF документ по пути к файлу и возвращает его структурное представление.
+
+ Args:
+ file_path (str): Путь к PDF файлу для парсинга.
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если файл не существует или не может быть прочитан.
+ """
+ logger.debug(f"Parsing PDF file: {file_path}")
+
+ if not os.path.exists(file_path):
+ raise ValueError(f"File not found: {file_path}")
+
+ try:
+ # Открываем PDF с помощью PyMuPDF
+ pdf_doc = fitz.open(file_path)
+ filename = os.path.basename(file_path)
+
+ return self._parse_document(pdf_doc, filename, file_path)
+
+ except Exception as e:
+ logger.error(f"Failed to open PDF file: {e}")
+ raise ValueError(f"Cannot open PDF file: {str(e)}")
+
+ def parse(self, file: BinaryIO, file_type: FileType | str | None = None) -> ParsedDocument:
+ """
+ Парсит PDF документ из объекта файла и возвращает его структурное представление.
+
+ Args:
+ file (BinaryIO): Объект файла для парсинга.
+ file_type: Тип файла, если известен.
+ Может быть объектом FileType или строкой с расширением (".pdf").
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если файл не может быть прочитан или распарсен.
+ """
+ logger.debug("Parsing PDF from file object")
+
+ # Проверяем соответствие типа файла
+ if file_type and isinstance(file_type, FileType) and file_type != FileType.PDF:
+ logger.warning(
+ f"Provided file_type {file_type} doesn't match parser type {FileType.PDF}"
+ )
+
+ try:
+ # Читаем содержимое файла в память
+ content = file.read()
+
+ # Открываем PDF из потока с помощью PyMuPDF
+ pdf_stream = io.BytesIO(content)
+ pdf_doc = fitz.open(stream=pdf_stream, filetype="pdf")
+
+ return self._parse_document(pdf_doc, "unknown.pdf", None)
+
+ except Exception as e:
+ logger.error(f"Failed to parse PDF from stream: {e}")
+ raise ValueError(f"Cannot parse PDF content: {str(e)}")
+
+ def _parse_document(
+ self,
+ pdf_doc: fitz.Document,
+ filename: str,
+ filepath: str | None,
+ ) -> ParsedDocument:
+ """
+ Внутренний метод для парсинга открытого PDF документа.
+
+ Args:
+ pdf_doc (fitz.Document): Открытый PDF документ.
+ filename (str): Имя файла для документа.
+ filepath (str | None): Путь к файлу (или None, если из объекта).
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если содержимое не может быть распарсено.
+ """
+ # Создание базового документа
+ doc = ParsedDocument(name=filename, type=FileType.PDF)
+
+ try:
+ # Извлечение метаданных
+ meta_dict = self.meta_parser.parse(pdf_doc, filepath)
+
+ # Преобразуем словарь метаданных в объект ParsedMeta
+ meta = ParsedMeta()
+ if 'author' in meta_dict:
+ meta.owner = meta_dict['author']
+ if 'creation_date' in meta_dict:
+ meta.date = meta_dict['creation_date']
+ if filepath:
+ meta.source = filepath
+
+ # Сохраняем остальные метаданные в поле note
+ meta.note = meta_dict
+
+ doc.meta = meta
+ logger.debug("Parsed metadata")
+
+ # Последовательный вызов парсеров
+ try:
+ # Парсим таблицы
+ doc.tables.extend(self.table_parser.parse(pdf_doc))
+ logger.debug(f"Parsed {len(doc.tables)} tables")
+
+ # Парсим изображения
+ doc.images.extend(self.image_parser.parse(pdf_doc))
+ logger.debug(f"Parsed {len(doc.images)} images")
+
+ # Парсим формулы
+ doc.formulas.extend(self.formula_parser.parse(pdf_doc))
+ logger.debug(f"Parsed {len(doc.formulas)} formulas")
+
+ # Парсим текст
+ doc.paragraphs.extend(self.paragraph_parser.parse(pdf_doc))
+ logger.debug(f"Parsed {len(doc.paragraphs)} paragraphs")
+
+ # Связываем элементы с их заголовками
+ self._link_elements_with_captions(doc)
+ logger.debug("Linked elements with captions")
+
+ except Exception as e:
+ logger.error(f"Error during parsing components: {e}")
+ logger.exception(e)
+ raise ValueError(f"Error parsing document components: {str(e)}")
+
+ return doc
+
+ finally:
+ # Закрываем документ после использования
+ pdf_doc.close()
+
+ def _link_elements_with_captions(self, doc: ParsedDocument) -> None:
+ """
+ Связывает таблицы, изображения и формулы с их заголовками на основе анализа текста.
+
+ Args:
+ doc (ParsedDocument): Документ для обработки.
+ """
+ # Находим параграфы, которые могут быть заголовками
+ caption_paragraphs = {}
+ for i, para in enumerate(doc.paragraphs):
+ text = para.text.lower()
+ if any(keyword in text for keyword in ["таблица", "рисунок", "формула", "рис.", "табл."]):
+ caption_paragraphs[i] = {
+ "text": text,
+ "page": para.page_number
+ }
+
+ # Для таблиц ищем соответствующие заголовки
+ for table in doc.tables:
+ table_page = table.page_number
+ # Ищем заголовки на той же странице или на предыдущей
+ for para_idx, caption_info in caption_paragraphs.items():
+ if ("таблица" in caption_info["text"] or "табл." in caption_info["text"]) and \
+ (caption_info["page"] == table_page or caption_info["page"] == table_page - 1):
+ table.title_index_in_paragraphs = para_idx
+ break
+
+ # Для изображений ищем соответствующие заголовки
+ for image in doc.images:
+ image_page = image.page_number
+ # Ищем заголовки на той же странице
+ for para_idx, caption_info in caption_paragraphs.items():
+ if ("рисунок" in caption_info["text"] or "рис." in caption_info["text"]) and \
+ caption_info["page"] == image_page:
+ image.referenced_element_index = para_idx
+ break
+
+ # Для формул ищем соответствующие заголовки
+ for formula in doc.formulas:
+ formula_page = formula.page_number
+ # Ищем заголовки на той же странице
+ for para_idx, caption_info in caption_paragraphs.items():
+ if "формула" in caption_info["text"] and caption_info["page"] == formula_page:
+ formula.referenced_element_index = para_idx
+ break
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/__init__.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2633670ef466b0cca567fdc3c06f9601965ba683
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/__init__.py
@@ -0,0 +1,15 @@
+"""
+Модуль для парсинга элементов XML-документов.
+"""
+
+from .formula_parser import XMLFormulaParser
+from .image_parser import XMLImageParser
+from .paragraph_parser import XMLParagraphParser
+from .table_parser import XMLTableParser
+
+__all__ = [
+ 'XMLTableParser',
+ 'XMLParagraphParser',
+ 'XMLImageParser',
+ 'XMLFormulaParser'
+]
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/formula_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/formula_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..2b188e318ebdeb9b64b6c3caed565d13aaf5f1b6
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/formula_parser.py
@@ -0,0 +1,57 @@
+"""
+Модуль для обработки формул из XML.
+"""
+
+from ....data_classes import ParsedFormula
+
+
+class XMLFormulaParser:
+ """
+ Класс для извлечения и обработки формул из XML документов.
+
+ Будет использовать BeautifulSoup для парсинга.
+ """
+
+ def parse(self, element) -> list[ParsedFormula]:
+ """
+ Извлекает формулу из XML элемента.
+
+ Args:
+ element: XML элемент формулы (будет использоваться BeautifulSoup Tag).
+
+ Returns:
+ Optional[ParsedFormula]: Объект с данными формулы или None, если формула не найдена.
+ """
+ # Создаем заглушку формул
+ formulas = []
+
+ # Полная реализация будет добавлена позже
+ # с использованием BeautifulSoup
+
+ return formulas
+
+ def extract_latex(self, element) -> str:
+ """
+ Извлекает LaTeX код формулы из XML элемента.
+
+ Args:
+ element: XML элемент.
+
+ Returns:
+ str: LaTeX код формулы.
+ """
+ # Заглушка для извлечения LaTeX кода
+ return ""
+
+ def extract_formula_metadata(self, element) -> dict[str, str]:
+ """
+ Извлекает метаданные формулы из XML элемента.
+
+ Args:
+ element: XML элемент.
+
+ Returns:
+ dict[str, str]: Словарь с метаданными формулы.
+ """
+ # Заглушка для извлечения метаданных
+ return {}
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/image_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/image_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..ba0586f74a200c710b33c23e71d5f8af32149c96
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/image_parser.py
@@ -0,0 +1,58 @@
+"""
+Модуль для обработки изображений из XML.
+"""
+
+from ....data_classes import ParsedImage
+
+
+class XMLImageParser:
+ """
+ Класс для извлечения и обработки изображений из XML документов.
+
+ Будет использовать BeautifulSoup для парсинга.
+ """
+
+ def parse(self, element) -> list[ParsedImage]:
+ """
+ Извлекает изображение из XML элемента.
+
+ Args:
+ element: XML элемент изображения (будет использоваться BeautifulSoup Tag).
+
+ Returns:
+ list[ParsedImage]: Список с данными изображений.
+ """
+ # Создаем заглушку изображения
+ images = []
+
+
+ # Полная реализация будет добавлена позже
+ # с использованием BeautifulSoup
+
+ return images
+
+ def extract_image_data(self, element) -> bytes:
+ """
+ Извлекает двоичные данные изображения из XML элемента.
+
+ Args:
+ element: XML элемент.
+
+ Returns:
+ bytes: Двоичные данные изображения.
+ """
+ # Заглушка для извлечения данных изображения
+ return bytes()
+
+ def extract_image_metadata(self, element) -> dict[str, str]:
+ """
+ Извлекает метаданные изображения из XML элемента.
+
+ Args:
+ element: XML элемент.
+
+ Returns:
+ dict[str, str]: Словарь с метаданными изображения.
+ """
+ # Заглушка для извлечения метаданных
+ return {}
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/meta_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/meta_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..9eab2d9622f1ed69392d5fe2e080b002154e6814
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/meta_parser.py
@@ -0,0 +1,99 @@
+"""
+Модуль для извлечения метаданных из XML документов.
+"""
+
+import logging
+import os
+from pathlib import Path
+
+from bs4 import BeautifulSoup
+
+from ....data_classes import ParsedMeta
+
+logger = logging.getLogger(__name__)
+
+
+class XMLMetaParser:
+ """
+ Класс для извлечения метаданных из XML документов.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует парсер метаданных.
+ """
+ # Теги для поиска информации о документе
+ self.name_tags = ["Название документа", "Наименование документа"]
+ self.owner_tags = ["Владелец процесса", "Владелец", "Ответственный"]
+ self.status_tags = ["Статус документа", "Статус"]
+
+ def parse(self, soup: BeautifulSoup, filepath: os.PathLike | None = None) -> ParsedMeta:
+ """
+ Извлекает метаданные из XML документа.
+
+ Args:
+ soup (BeautifulSoup): Объект BeautifulSoup с XML документом.
+ filepath (os.PathLike | None): Путь к файлу (для извлечения имени файла).
+
+ Returns:
+ ParsedMeta: Объект с метаданными документа.
+ """
+ # Извлекаем базовую информацию
+ name = self._extract_info_recurse(soup, self.name_tags)
+ owner = self._extract_info_recurse(soup, self.owner_tags)
+ status = self._extract_info_recurse(soup, self.status_tags)
+
+ # Если не удалось извлечь имя, используем имя файла
+ if not name and filepath:
+ name = Path(filepath).stem
+
+ # Если не удалось извлечь владельца, используем дефолтное значение
+ if not owner:
+ owner = "-"
+
+ logger.debug(f"Extracted metadata - Name: {name}, Owner: {owner}, Status: {status}")
+
+ return ParsedMeta(
+ owner=owner,
+ status=status,
+ )
+
+ def _extract_info_value(self, soup: BeautifulSoup, key: str) -> str:
+ """
+ Извлекает значение по ключевому тегу.
+
+ Args:
+ soup (BeautifulSoup): Объект BeautifulSoup с XML документом.
+ key (str): Ключевой тег для поиска.
+
+ Returns:
+ str: Найденное значение или пустая строка.
+ """
+ # Ищем тег в документе
+ key_tag = soup.find(string=key)
+ if not key_tag:
+ return ''
+
+ # Ищем следующий текстовый тег после ключевого
+ next_tag = key_tag.find_next('w:t')
+ if not next_tag:
+ return ''
+
+ return next_tag.get_text().strip()
+
+ def _extract_info_recurse(self, soup: BeautifulSoup, keys: list[str]) -> str:
+ """
+ Извлекает значение, перебирая несколько возможных ключевых тегов.
+
+ Args:
+ soup (BeautifulSoup): Объект BeautifulSoup с XML документом.
+ keys (list[str]): Список ключевых тегов для поиска.
+
+ Returns:
+ str: Первое найденное значение или пустая строка.
+ """
+ for key in keys:
+ value = self._extract_info_value(soup, key)
+ if value:
+ return value
+ return ''
\ No newline at end of file
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/paragraph_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/paragraph_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..a544cd45e3e525c3c35eb014fa84d0101f598314
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/paragraph_parser.py
@@ -0,0 +1,611 @@
+"""
+Модуль для обработки текстовых блоков (параграфов) из XML.
+"""
+
+import logging
+import re
+from typing import Any, List
+
+from bs4 import BeautifulSoup, Tag
+
+from ....data_classes import ParsedTextBlock, TextStyle
+
+logger = logging.getLogger(__name__)
+
+
+class XMLParagraphParser:
+ """
+ Класс для извлечения и обработки текстовых блоков из XML документов.
+
+ Примечание: Этот парсер модифицирует переданный soup объект, удаляя таблицы,
+ бинарные данные, изображения и др. Поэтому его следует вызывать последним
+ в цепочке парсеров.
+ """
+
+ def __init__(
+ self,
+ style_cache: dict[str, Any] | None = None,
+ numbering_cache: dict[str, Any] | None = None,
+ relationships_cache: dict[str, Any] | None = None
+ ):
+ """
+ Инициализирует парсер параграфов.
+
+ Args:
+ style_cache (dict[str, Any] | None): Кэш стилей из DocxParser
+ numbering_cache (dict[str, Any] | None): Кэш нумерации из DocxParser
+ relationships_cache (dict[str, Any] | None): Кэш связей из DocxParser
+ """
+ self.style_cache = style_cache or {}
+ self.numbering_cache = numbering_cache or {}
+ self.relationships_cache = relationships_cache or {}
+ # Сохраняем информацию о закладках
+ self.bookmarks = {} # имя -> (начало, конец)
+
+ # Включаем подробное логирование для отладки
+ self.debug = True
+
+ def parse(self, soup: BeautifulSoup) -> List[ParsedTextBlock]:
+ """
+ Извлекает все текстовые блоки из XML документа.
+
+ Args:
+ soup (BeautifulSoup): Объект BeautifulSoup с XML документом.
+ Внимание: этот метод изменяет переданный soup объект!
+
+ Returns:
+ List[ParsedTextBlock]: Список извлеченных текстовых блоков.
+ """
+ # Удаляем элементы, которые не должны быть в тексте
+ # Внимание: это изменяет переданный soup объект
+ self._remove_non_text_elements(soup)
+
+ # Предварительный проход для сбора всех закладок
+ self._collect_bookmarks(soup)
+
+ # Собираем сноски
+ footnotes = self._collect_footnotes(soup)
+ logger.debug(f"Collected {len(footnotes)} footnotes")
+
+ paragraphs = []
+ # Извлекаем абзацы (теги w:p)
+ for p_tag in soup.find_all('w:p'):
+ paragraph = self.parse_paragraph(p_tag)
+
+ # Добавляем сноски, если они есть в этом параграфе
+ paragraph_footnotes = self._get_paragraph_footnotes(p_tag, footnotes)
+ if paragraph_footnotes:
+ paragraph.footnotes.extend(paragraph_footnotes)
+
+ if paragraph.text.strip(): # Добавляем только непустые абзацы
+ paragraphs.append(paragraph)
+
+ logger.debug(f"Extracted {len(paragraphs)} non-empty paragraphs")
+ return paragraphs
+
+ def _collect_bookmarks(self, soup: BeautifulSoup) -> None:
+ """
+ Собирает информацию о закладках в документе.
+
+ Args:
+ soup (BeautifulSoup): Объект BeautifulSoup с XML документом.
+ """
+ bookmarks_start = {}
+ bookmarks_end = {}
+
+ # Собираем все стартовые теги закладок
+ for bookmark_start in soup.find_all('w:bookmarkStart'):
+ if 'w:id' in bookmark_start.attrs and 'w:name' in bookmark_start.attrs:
+ bookmark_id = bookmark_start['w:id']
+ bookmark_name = bookmark_start['w:name']
+ bookmarks_start[bookmark_id] = {'name': bookmark_name, 'element': bookmark_start}
+
+ # Собираем все конечные теги закладок
+ for bookmark_end in soup.find_all('w:bookmarkEnd'):
+ if 'w:id' in bookmark_end.attrs and bookmark_end['w:id'] in bookmarks_start:
+ bookmark_id = bookmark_end['w:id']
+ bookmarks_end[bookmark_id] = {'element': bookmark_end}
+
+ if self.debug:
+ logger.debug(f"Found {len(bookmarks_start)} bookmarks in document")
+
+ # Сохраняем информацию о закладках для использования при обработке параграфов
+ self.bookmarks = {
+ bookmarks_start[bookmark_id]['name']: {
+ 'id': bookmark_id,
+ 'start': bookmarks_start[bookmark_id]['element'],
+ 'end': bookmarks_end.get(bookmark_id, {}).get('element')
+ }
+ for bookmark_id in bookmarks_start
+ if bookmark_id in bookmarks_end
+ }
+
+ def _remove_non_text_elements(self, soup: BeautifulSoup) -> None:
+ """
+ Удаляет нетекстовые элементы из документа.
+
+ Args:
+ soup (BeautifulSoup): Объект BeautifulSoup с XML документом.
+ """
+ # Удаляем таблицы
+ for tbl in soup.find_all('w:tbl'):
+ tbl.decompose()
+
+ # Удаляем бинарные данные
+ for binary in soup.find_all('w:binData'):
+ binary.decompose()
+
+ # Удаляем изображения
+ for drawing in soup.find_all('w:drawing'):
+ drawing.decompose()
+
+ # Удаляем объекты
+ for object_tag in soup.find_all('w:object'):
+ object_tag.decompose()
+
+ # Удаляем комментарии
+ for comment in soup.find_all('w:commentRangeStart'):
+ comment.decompose()
+
+ for comment in soup.find_all('w:commentRangeEnd'):
+ comment.decompose()
+
+ for comment in soup.find_all('w:commentReference'):
+ comment.decompose()
+
+ # НЕ удаляем сноски, а собираем их для дальнейшей обработки
+
+ def parse_paragraph(self, element: Tag) -> ParsedTextBlock:
+ """
+ Извлекает текстовый блок из XML элемента.
+
+ Args:
+ element (Tag): XML элемент параграфа.
+
+ Returns:
+ ParsedTextBlock: Объект с данными текстового блока.
+ """
+ text = ""
+ runs = []
+
+ # Получаем стилевую информацию и информацию о нумерации
+ style_info = self.extract_paragraph_style(element)
+
+ # Статистика для распознавания форматирования
+ total_runs = 0
+ bold_runs = 0
+ italic_runs = 0
+ underline_runs = 0
+
+ # Сбор идентификаторов закладок в параграфе (якорей)
+ anchors = set() # Просто множество строк
+ for bookmark_start in element.find_all('w:bookmarkStart'):
+ if 'w:id' in bookmark_start.attrs and 'w:name' in bookmark_start.attrs:
+ bookmark_name = bookmark_start['w:name']
+ anchors.add(bookmark_name)
+ if self.debug:
+ logger.debug(f"Found bookmark (anchor): {bookmark_name}")
+
+ # Сбор идентификаторов ссылок в параграфе
+ links = set() # Просто множество строк
+
+ # 1. Обычные гиперссылки
+ for hyperlink in element.find_all('w:hyperlink'):
+ link_target = ""
+
+ if 'r:id' in hyperlink.attrs:
+ rel_id = hyperlink['r:id']
+ if rel_id in self.relationships_cache:
+ rel_info = self.relationships_cache[rel_id]
+ if rel_info['type'] == 'hyperlink':
+ link_target = rel_info['target']
+ elif rel_info['type'] == 'bookmark':
+ link_target = f"#bookmark:{rel_info['target']}"
+ elif 'w:anchor' in hyperlink.attrs:
+ link_target = f"#anchor:{hyperlink['w:anchor']}"
+ elif 'w:bookmark' in hyperlink.attrs:
+ link_target = f"#bookmark:{hyperlink['w:bookmark']}"
+ elif 'w:bookmarkRef' in hyperlink.attrs:
+ link_target = f"#bookmark:{hyperlink['w:bookmarkRef']}"
+ elif 'w:dest' in hyperlink.attrs:
+ link_target = hyperlink['w:dest']
+
+ if link_target:
+ links.add(link_target)
+ if self.debug:
+ logger.debug(f"Found hyperlink: {link_target}")
+
+ # 2. Перекрестные ссылки - простые поля
+ for fld_simple in element.find_all('w:fldSimple'):
+ if 'w:instr' in fld_simple.attrs:
+ instr = fld_simple['w:instr']
+ target = self._parse_field_instruction(instr)
+
+ if target:
+ links.add(target)
+ if self.debug:
+ logger.debug(f"Found simple cross-reference to: {target}")
+
+ # 3. Перекрестные ссылки - сложные поля
+ inside_field = False
+ current_instr = ""
+
+ for run in element.find_all('w:r'):
+ fld_char = run.find('w:fldChar')
+ instr_text = run.find('w:instrText')
+
+ if fld_char and 'w:fldCharType' in fld_char.attrs:
+ fld_type = fld_char['w:fldCharType']
+
+ if fld_type == 'begin':
+ inside_field = True
+ current_instr = ""
+ elif fld_type == 'end' and inside_field:
+ inside_field = False
+ target = self._parse_field_instruction(current_instr)
+ if target:
+ links.add(target)
+ if self.debug:
+ logger.debug(f"Found complex cross-reference to: {target}")
+
+ if inside_field and instr_text:
+ current_instr += instr_text.get_text().strip()
+
+ # Обрабатываем каждый текстовый запуск (run)
+ for run in element.find_all('w:r'):
+ run_text = self._extract_run_text(run)
+
+ # Если у запуска есть текст
+ if run_text:
+ total_runs += 1
+
+ # Получаем информацию о стиле
+ run_style = self._extract_minimal_run_style(run)
+
+ # Обновляем статистику форматирования
+ if run_style.get('bold'):
+ bold_runs += 1
+ if run_style.get('italic'):
+ italic_runs += 1
+ if run_style.get('underline'):
+ underline_runs += 1
+
+ # Добавляем в runs, только если есть стилевое форматирование
+ if run_style:
+ runs.append({"text": run_text, "style": run_style})
+
+ # Добавляем текст к общему тексту параграфа
+ text += run_text
+
+ # Очищаем объединенный текст
+ text = self._clean_text(text)
+
+ # Определяем fully/partly для форматирования
+ style = TextStyle()
+
+ # Заполняем информацию о стиле параграфа
+ style.paragraph_style = style_info.get('paragraph_style', '')
+ style.alignment = style_info.get('alignment', '')
+
+ # Заполняем информацию о нумерации
+ if 'numbering' in style_info:
+ numbering_info = style_info['numbering']
+ style.has_numbering = True
+ style.numbering_id = numbering_info.get('id', '')
+ style.numbering_level = int(numbering_info.get('level', '0'))
+
+ # Получаем формат нумерации из кэша
+ if self.numbering_cache and style.numbering_id in self.numbering_cache:
+ num_info = self.numbering_cache[style.numbering_id]
+ if 'levels' in num_info and str(style.numbering_level) in num_info['levels']:
+ level_info = num_info['levels'][str(style.numbering_level)]
+ style.numbering_format = level_info.get('format', '')
+
+ # Преобразуем ID стиля в имя стиля из кэша
+ if style.paragraph_style and self.style_cache:
+ style_id = style.paragraph_style
+ if style_id in self.style_cache:
+ style.paragraph_style_name = self.style_cache[style_id].get('name', '')
+
+ # Устанавливаем флаги форматирования на основе статистики
+ if total_runs > 0:
+ # Bold
+ if bold_runs == total_runs:
+ style.fully_bold = True
+ elif bold_runs > 0:
+ style.partly_bold = True
+
+ # Italic
+ if italic_runs == total_runs:
+ style.fully_italic = True
+ elif italic_runs > 0:
+ style.partly_italic = True
+
+ # Underline
+ if underline_runs == total_runs:
+ style.fully_underlined = True
+ elif underline_runs > 0:
+ style.partly_underlined = True
+
+ # Создаем текстовый блок
+ parsed_block = ParsedTextBlock(
+ text=text,
+ style=style,
+ links=list(links), # Преобразуем множество в список строк
+ anchors=list(anchors), # Преобразуем множество в список строк
+ metadata=runs # Храним только runs с непустыми стилями
+ )
+
+ if self.debug:
+ logger.debug(f"Created parsed block with {len(links)} links and {len(anchors)} anchors")
+ logger.debug(f"Text: '{text[:100]}{'...' if len(text) > 100 else ''}'")
+
+ return parsed_block
+
+ def _extract_minimal_run_style(self, run: Tag) -> dict:
+ """
+ Извлекает только жирность, курсив и подчеркивание из запуска текста.
+
+ Args:
+ run (Tag): XML элемент запуска текста.
+
+ Returns:
+ dict: Словарь с информацией о форматировании или пустой словарь, если
+ нет форматирования.
+ """
+ style_info = {}
+ r_pr = run.find('w:rPr')
+ if r_pr:
+ # Жирный
+ b_tag = r_pr.find('w:b')
+ if b_tag:
+ is_bold = True
+ if 'w:val' in b_tag.attrs and b_tag['w:val'] == '0':
+ is_bold = False
+ if is_bold: # Добавляем только если действительно жирный
+ style_info['bold'] = True
+
+ # Курсив
+ i_tag = r_pr.find('w:i')
+ if i_tag:
+ is_italic = True
+ if 'w:val' in i_tag.attrs and i_tag['w:val'] == '0':
+ is_italic = False
+ if is_italic: # Добавляем только если действительно курсив
+ style_info['italic'] = True
+
+ # Подчеркнутый
+ u_tag = r_pr.find('w:u')
+ if u_tag:
+ if 'w:val' in u_tag.attrs and u_tag['w:val'] != 'none':
+ style_info['underline'] = True
+
+ return style_info
+
+ def _clean_text(self, text: str) -> str:
+ """
+ Очищает текст от специальных символов и ненужных элементов.
+
+ Args:
+ text (str): Исходный текст.
+
+ Returns:
+ str: Очищенный текст.
+ """
+ # Замена HTML-сущностей
+ text = text.replace('&', '&')
+ text = text.replace('<', '<')
+ text = text.replace('>', '>')
+
+ # Удаление специфических строк
+ text = text.replace('MS-Word', '')
+ text = text.replace('См. документ в ', '')
+ text = text.replace(
+ '------------------------------------------------------------------', ''
+ )
+
+ # Удаление фигурных скобок и их содержимого
+ text = re.sub(
+ r'\{[\.\,\#\:\=A-Za-zа-яА-Я\d\/\s\"\-\/\?\%\_\.\&\$]+\}', '', text
+ )
+
+ return text.strip()
+
+ def extract_links(self, element: Tag) -> List[str]:
+ """
+ Извлекает идентификаторы ссылок из XML элемента.
+ УСТАРЕВШИЙ МЕТОД: вместо него теперь используется полный анализ ссылок в parse_paragraph.
+
+ Args:
+ element (Tag): XML элемент.
+
+ Returns:
+ List[str]: Список идентификаторов извлеченных ссылок.
+ """
+ links = []
+ # Ищем гиперссылки в элементе
+ for hyperlink in element.find_all('w:hyperlink'):
+ if 'r:id' in hyperlink.attrs:
+ links.append(hyperlink['r:id'])
+ return links
+
+ def extract_paragraph_style(self, element: Tag) -> dict:
+ """
+ Извлекает стиль параграфа, выравнивание и информацию о нумерации из XML элемента.
+
+ Args:
+ element (Tag): XML элемент.
+
+ Returns:
+ dict: Словарь с информацией о стиле параграфа, выравнивании и нумерации.
+ """
+ style_info = {}
+
+ # Извлекаем информацию о стиле параграфа
+ p_pr = element.find('w:pPr')
+ if p_pr:
+ # Стиль параграфа
+ p_style = p_pr.find('w:pStyle')
+ if p_style and 'w:val' in p_style.attrs:
+ style_info['paragraph_style'] = p_style['w:val']
+
+ # Выравнивание
+ jc = p_pr.find('w:jc')
+ if jc and 'w:val' in jc.attrs:
+ style_info['alignment'] = jc['w:val']
+
+ # Нумерация
+ num_pr = p_pr.find('w:numPr')
+ if num_pr:
+ numbering = {}
+ ilvl = num_pr.find('w:ilvl')
+ if ilvl and 'w:val' in ilvl.attrs:
+ numbering['level'] = ilvl['w:val']
+ num_id = num_pr.find('w:numId')
+ if num_id and 'w:val' in num_id.attrs:
+ numbering['id'] = num_id['w:val']
+ if numbering:
+ style_info['numbering'] = numbering
+
+ return style_info
+
+ def _extract_run_text(self, run: Tag) -> str:
+ """
+ Извлекает текст из w:r элемента.
+
+ Args:
+ run (Tag): XML элемент w:r
+
+ Returns:
+ str: Извлеченный текст
+ """
+ run_text = ""
+ for text_tag in run.find_all('w:t'):
+ content = text_tag.get_text()
+
+ # Пропускаем специальные элементы в фигурных скобках
+ if content and '{' in content and '}' in content:
+ if '{КСС}' in content or content.startswith('{СС_'):
+ continue
+
+ if content:
+ run_text += content
+ return run_text
+
+ def _parse_field_instruction(self, instr: str) -> str | None:
+ """
+ Разбирает инструкцию поля и извлекает цель ссылки.
+
+ Args:
+ instr (str): Инструкция поля.
+
+ Returns:
+ str | None: Цель ссылки или None, если не найдена.
+ """
+ if not instr:
+ return None
+
+ # Нормализуем пробелы и приводим к нижнему регистру для упрощения парсинга
+ instr = ' '.join(instr.split()).lower()
+
+ # REF - перекрестная ссылка
+ if instr.startswith('ref'):
+ # Извлекаем имя закладки
+ match = re.search(r'ref\s+([^\s\\]+)', instr)
+ if match:
+ bookmark_name = match.group(1)
+ return f"#bookmark:{bookmark_name}"
+
+ # PAGEREF - ссылка на страницу
+ elif instr.startswith('pageref'):
+ match = re.search(r'pageref\s+([^\s\\]+)', instr)
+ if match:
+ bookmark_name = match.group(1)
+ return f"#page:{bookmark_name}"
+
+ # HYPERLINK - прямая гиперссылка
+ elif instr.startswith('hyperlink'):
+ match = re.search(r'hyperlink\s+"([^"]+)"', instr)
+ if match:
+ url = match.group(1)
+ return url
+
+ return None
+
+ def _collect_footnotes(self, soup: BeautifulSoup) -> dict[str, dict[str, Any]]:
+ """
+ Собирает все сноски из документа.
+
+ Args:
+ soup (BeautifulSoup): Объект BeautifulSoup с XML документом.
+
+ Returns:
+ dict[str, dict[str, Any]]: Словарь сносок с id в качестве ключа.
+ """
+ footnotes = {}
+
+ # Ищем все ссылки на сноски
+ for footnote_ref in soup.find_all('w:footnoteReference'):
+ if 'w:id' in footnote_ref.attrs:
+ footnote_id = footnote_ref['w:id']
+ footnotes[footnote_id] = {"id": footnote_id, "content": "", "reference": footnote_ref}
+
+ # Ищем все сноски в документе (обычно в отдельном разделе)
+ for footnote in soup.find_all('w:footnote'):
+ if 'w:id' in footnote.attrs and footnote['w:id'] in footnotes:
+ footnote_id = footnote['w:id']
+ footnote_text = ""
+
+ # Извлекаем текст сноски
+ for p in footnote.find_all('w:p'):
+ for t in p.find_all('w:t'):
+ footnote_text += t.get_text()
+
+ footnotes[footnote_id]["content"] = footnote_text.strip()
+
+ # Удаляем сноски без контента
+ footnotes = {k: v for k, v in footnotes.items() if v["content"]}
+
+ if self.debug:
+ logger.debug(f"Found {len(footnotes)} footnotes in document")
+
+ return footnotes
+
+ def _get_paragraph_footnotes(self, paragraph: Tag, footnotes: dict[str, dict[str, Any]]) -> list[dict[str, Any]]:
+ """
+ Получает список сносок, относящихся к данному параграфу.
+
+ Args:
+ paragraph (Tag): XML элемент параграфа.
+ footnotes (dict[str, dict[str, Any]]): Словарь всех сносок документа.
+
+ Returns:
+ list[dict[str, Any]]: Список сносок параграфа.
+ """
+ paragraph_footnotes = []
+
+ # Ищем ссылки на сноски в параграфе
+ for footnote_ref in paragraph.find_all('w:footnoteReference'):
+ if 'w:id' in footnote_ref.attrs and footnote_ref['w:id'] in footnotes:
+ footnote_id = footnote_ref['w:id']
+ footnote_info = footnotes[footnote_id]
+
+ # Создаем объект сноски для параграфа
+ footnote_obj = {
+ "id": footnote_id,
+ "text": footnote_info["content"],
+ "marker": footnote_id # Используем ID как маркер, если нет другого
+ }
+
+ # Ищем маркер сноски (обычно это цифра в тексте перед ссылкой)
+ run = footnote_ref.parent
+ if run and run.name == 'w:r':
+ prev_run = run.find_previous_sibling('w:r')
+ if prev_run:
+ text_ele = prev_run.find('w:t')
+ if text_ele and text_ele.get_text().strip().isdigit():
+ footnote_obj["marker"] = text_ele.get_text().strip()
+
+ paragraph_footnotes.append(footnote_obj)
+
+ return paragraph_footnotes
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/table/__init__.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/table/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/table_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/table_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b4c770bdaa3629046d134e05a9148d73580a606
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml/table_parser.py
@@ -0,0 +1,723 @@
+"""
+Модуль для обработки таблиц из XML.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+from bs4 import BeautifulSoup, Tag
+
+from ....data_classes import ParsedRow, ParsedSubtable, ParsedTable, TextStyle
+from .paragraph_parser import XMLParagraphParser
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TableStyle:
+ """
+ Класс для представления стиля таблицы.
+ """
+ style_id: str = ""
+ style_name: str = ""
+ width: str = ""
+ width_type: str = ""
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Конвертирует стиль таблицы в словарь для обратной совместимости.
+
+ Returns:
+ dict[str, Any]: Словарное представление стиля
+ """
+ result = {}
+ if self.style_id:
+ result['style_id'] = self.style_id
+ if self.style_name:
+ result['style_name'] = self.style_name
+ if self.width:
+ result['width'] = self.width
+ if self.width_type:
+ result['width_type'] = self.width_type
+ return result
+
+
+@dataclass
+class TableFormatStats:
+ """
+ Класс для хранения статистики форматирования таблицы.
+ Используется для определения заголовков и особенностей таблицы.
+ """
+ total_rows: int = 0
+ bold_rows: int = 0
+ italic_rows: int = 0
+ center_aligned_rows: int = 0
+ numeric_content_rows: int = 0
+ expected_cells: int = 0
+ grid_cols: list[int] = field(default_factory=list) # Ширины столбцов из сетки таблицы
+
+
+@dataclass
+class RowInfo:
+ """
+ Класс для хранения информации о строке таблицы и её роли.
+ """
+ row_element: Tag # XML элемент строки
+ parsed_row: ParsedRow # Распарсенная строка
+ is_header: bool = False # Является ли строка заголовком
+ is_subtitle: bool = False # Является ли строка подзаголовком
+ is_note: bool = False # Является ли строка примечанием
+
+
+class XMLTableParser:
+ """
+ Класс для извлечения и обработки таблиц из XML документов.
+
+ Использует BeautifulSoup для парсинга XML-структуры таблиц.
+ Обрабатывает объединённые ячейки, подтаблицы и извлекает форматирование.
+ """
+
+ def __init__(self, style_cache: dict[str, Any] | None = None):
+ """
+ Инициализирует парсер таблиц.
+
+ Args:
+ style_cache (dict[str, Any] | None): Кэш стилей из DocxParser
+ """
+ self.style_cache = style_cache or {}
+ self.paragraph_parser = XMLParagraphParser(style_cache=style_cache)
+
+ def parse(self, soup: BeautifulSoup) -> list[ParsedTable]:
+ """
+ Извлекает все таблицы из XML документа.
+
+ Args:
+ soup (BeautifulSoup): Объект BeautifulSoup с XML документом.
+
+ Returns:
+ list[ParsedTable]: Список извлеченных таблиц.
+ """
+ tables = []
+ table_elements = soup.find_all('w:tbl')
+
+ logger.debug(f"Найдено {len(table_elements)} таблиц в документе")
+
+ for i, table_element in enumerate(table_elements):
+ # Получаем заголовок таблицы из предыдущего параграфа, если он есть
+ title = self._extract_table_title(table_element) # работает
+
+ # Получаем информацию о сетке таблицы
+ grid_cols = self._get_table_grid(table_element)
+
+ # Извлекаем все строки таблицы и классифицируем их
+ rows_info = self._extract_and_classify_rows(table_element, grid_cols)
+
+ # Если таблица пустая, пропускаем её
+ if not rows_info:
+ continue
+
+ # Определяем заголовок таблицы
+ headers = self._extract_headers(rows_info)
+
+ # Определяем примечание (последняя строка, которая помечена как примечание)
+ note = None
+ for row_info in reversed(rows_info):
+ if row_info.is_note:
+ # Берем текст из первой ячейки строки примечания
+ note = row_info.parsed_row.cells[0].strip()
+ # Удаляем строку примечания из списка строк
+ rows_info.remove(row_info)
+ break
+
+ # Разбиваем строки на подтаблицы по подзаголовкам
+ subtables = self._split_into_subtables(rows_info)
+
+ # Нормализуем ширину строк
+ expected_width = len(grid_cols) or (len(headers[0].cells) if headers else 0)
+ if expected_width == 0 and subtables:
+ # Если нет информации о ширине, используем ширину первой строки первой подтаблицы
+ for subtable in subtables:
+ if subtable.rows:
+ expected_width = len(subtable.rows[0].cells)
+ break
+
+ self._normalize_row_widths(headers, subtables, expected_width)
+
+ # Создаем итоговую таблицу
+ parsed_table = ParsedTable(
+ title=title,
+ note=note,
+ index=[str(i + 1)],
+ headers=headers,
+ subtables=subtables,
+ )
+
+ tables.append(parsed_table)
+
+ logger.debug(f"Извлечено {len(tables)} непустых таблиц")
+ return tables
+
+ def _get_table_grid(self, table_element: Tag) -> list[int]:
+ """
+ Извлекает информацию о сетке таблицы.
+
+ Args:
+ table_element (Tag): XML элемент таблицы
+
+ Returns:
+ list[int]: Список ширин столбцов
+ """
+ grid = table_element.find('w:tblGrid')
+ if not grid:
+ return []
+
+ grid_cols = grid.find_all('w:gridCol')
+ widths = []
+
+ for col in grid_cols:
+ if 'w:w' in col.attrs:
+ widths.append(int(col['w:w']))
+ else:
+ widths.append(0)
+
+ return widths
+
+ def _extract_table_style(self, table_element: Tag) -> TableStyle:
+ """
+ Извлекает информацию о стиле таблицы.
+
+ Args:
+ table_element (Tag): XML элемент таблицы
+
+ Returns:
+ TableStyle: Стиль таблицы
+ """
+ style = TableStyle()
+
+ # Ищем свойства таблицы
+ tbl_pr = table_element.find('w:tblPr')
+ if tbl_pr:
+ # Стиль таблицы
+ tbl_style = tbl_pr.find('w:tblStyle')
+ if tbl_style and 'w:val' in tbl_style.attrs:
+ style.style_id = tbl_style['w:val']
+
+ # Если есть кэш стилей, получаем имя стиля
+ if self.style_cache and style.style_id in self.style_cache:
+ style.style_name = self.style_cache[style.style_id].get('name', '')
+
+ # Ширина таблицы
+ tbl_w = tbl_pr.find('w:tblW')
+ if tbl_w:
+ if 'w:w' in tbl_w.attrs:
+ style.width = tbl_w['w:w']
+ if 'w:type' in tbl_w.attrs:
+ style.width_type = tbl_w['w:type']
+
+ return style
+
+ def _extract_table_title(self, table_element: Tag) -> str | None:
+ """
+ Извлекает заголовок таблицы из предшествующего параграфа.
+
+ Args:
+ table_element (Tag): XML элемент таблицы
+
+ Returns:
+ str | None: Заголовок таблицы или None, если не найден
+ """
+ # Ищем последний предшествующий параграф
+ previous_paragraph = table_element.find_previous('w:p')
+ if previous_paragraph:
+ parsed = self.paragraph_parser.parse_paragraph(previous_paragraph)
+ if parsed.text:
+ return parsed.text
+
+ return None
+
+ def _extract_and_classify_rows(self, table_element: Tag, grid_cols: list[int]) -> list[RowInfo]:
+ """
+ Извлекает все строки таблицы и классифицирует их (заголовок, подзаголовок, примечание).
+
+ Args:
+ table_element (Tag): XML элемент таблицы
+ grid_cols (list[int]): Информация о сетке таблицы
+
+ Returns:
+ list[RowInfo]: Список информации о строках таблицы
+ """
+ rows_info = []
+ row_elements = table_element.find_all('w:tr')
+
+ # Определяем ожидаемую ширину таблицы на основе сетки
+ expected_width = len(grid_cols)
+
+ # Первый проход - извлекаем все строки и сразу определяем, является ли строка
+ # подзаголовком или примечанием (до дублирования ячеек)
+ for row_index, row_element in enumerate(row_elements):
+ # Парсим строку
+ parsed_row, real_cells = self._parse_row(row_element)
+
+ # Создаем объект с информацией о строке
+ row_info = RowInfo(
+ row_element=row_element,
+ parsed_row=parsed_row
+ )
+
+ # Определяем, является ли строка подзаголовком или примечанием
+ # Проверяем наличие объединенной ячейки на всю ширину таблицы
+ if self._is_subtitle_row(row_element, parsed_row, expected_width, real_cells):
+ row_info.is_subtitle = True
+
+ # Проверяем, является ли строка примечанием (последняя строка с объединенной ячейкой)
+ if row_index == len(row_elements) - 1 and self._is_note_row(row_element, parsed_row, expected_width, real_cells):
+ row_info.is_note = True
+
+ rows_info.append(row_info)
+
+ return rows_info
+
+ def _extract_headers(self, rows_info: list[RowInfo]) -> list[ParsedRow]:
+ """
+ Определяет, какие строки являются заголовками таблицы.
+
+ Args:
+ rows_info (list[RowInfo]): Список информации о строках таблицы
+
+ Returns:
+ list[ParsedRow]: Список строк заголовка
+ """
+ if not rows_info:
+ return []
+
+ headers = []
+
+ # Собираем статистику по форматированию строк
+ format_stats = self._analyze_rows_formatting(rows_info)
+
+ # Помечаем строки как заголовки на основе форматирования
+ header_rows_indices = self._identify_header_rows(rows_info, format_stats)
+
+ # Выбираем строки заголовка
+ for i in header_rows_indices:
+ if i < len(rows_info):
+ rows_info[i].is_header = True
+ # Создаем копию строки с флагом заголовка
+ header_row = ParsedRow(
+ index=rows_info[i].parsed_row.index,
+ cells=rows_info[i].parsed_row.cells.copy(),
+ style=rows_info[i].parsed_row.style,
+ is_header=True
+ )
+ headers.append(header_row)
+
+ return headers
+
+ def _analyze_rows_formatting(self, rows_info: list[RowInfo]) -> TableFormatStats:
+ """
+ Анализирует форматирование всех строк для определения общих характеристик.
+
+ Args:
+ rows_info (list[RowInfo]): Список информации о строках таблицы
+
+ Returns:
+ TableFormatStats: Статистика форматирования таблицы
+ """
+ stats = TableFormatStats(total_rows=len(rows_info))
+
+ # Подсчитываем строки с различным форматированием
+ for row_info in rows_info:
+ row = row_info.parsed_row
+
+ if row.style.fully_bold:
+ stats.bold_rows += 1
+ if row.style.fully_italic:
+ stats.italic_rows += 1
+ if row.style.alignment == 'center':
+ stats.center_aligned_rows += 1
+
+ # Проверяем, содержат ли ячейки преимущественно числовые данные
+ numeric_cells = 0
+ for cell in row.cells:
+ # Проверка на числовые данные (включая числа с разделителями)
+ if cell.strip() and all(c.isdigit() or c in '., -/+%' for c in cell.strip()):
+ numeric_cells += 1
+
+ if numeric_cells > len(row.cells) / 2: # Если больше половины ячеек - числа
+ stats.numeric_content_rows += 1
+
+ return stats
+
+ def _identify_header_rows(self, rows_info: list[RowInfo], format_stats: TableFormatStats) -> list[int]:
+ """
+ Определяет индексы строк, которые являются заголовками.
+
+ Args:
+ rows_info (list[RowInfo]): Список информации о строках таблицы
+ format_stats (TableFormatStats): Статистика форматирования
+
+ Returns:
+ list[int]: Индексы строк-заголовков
+ """
+ if not rows_info:
+ return []
+
+ header_indices = []
+ total_rows = format_stats.total_rows
+
+ # Определяем, какие критерии форматирования не являются всеобщими
+ # и могут использоваться для идентификации заголовков
+ is_bold_unique = format_stats.bold_rows < total_rows
+ is_italic_unique = format_stats.italic_rows < total_rows
+ is_center_unique = format_stats.center_aligned_rows < total_rows
+
+ # Проверяем первые строки таблицы (до 3-4 строк)
+ for i, row_info in enumerate(rows_info[:4]):
+ # Пропускаем строки, которые уже определены как подзаголовки или примечания
+ if row_info.is_subtitle or row_info.is_note:
+ continue
+
+ row = row_info.parsed_row
+ row_tag = row_info.row_element
+
+ # Проверяем, может ли строка быть заголовком
+ is_header = False
+
+ # Проверяем наличие специального атрибута заголовка
+ tr_pr = row_tag.find('w:trPr')
+ if tr_pr and tr_pr.find('w:tblHeader'):
+ is_header = True
+ elif is_bold_unique and row.style.fully_bold:
+ is_header = True
+ elif is_italic_unique and row.style.fully_italic:
+ is_header = True
+ elif is_center_unique and row.style.alignment == 'center':
+ is_header = True
+
+ # Дополнительно проверяем первую строку - она часто является заголовком
+ # даже без особого форматирования
+ if i == 0 and not header_indices:
+ is_header = True
+
+ # Проверяем, содержит ли строка преимущественно числа
+ # Если да, то это скорее всего НЕ заголовок
+ numeric_cells = 0
+ for cell in row.cells:
+ if cell.strip() and all(c.isdigit() or c in '., ' for c in cell.strip()):
+ numeric_cells += 1
+
+ if numeric_cells > len(row.cells) / 2:
+ is_header = False
+
+ # Если строка определена как заголовок, добавляем её индекс
+ if is_header:
+ header_indices.append(i)
+ else:
+ # Если текущая строка не заголовок, прекращаем поиск заголовков
+ break
+
+ return header_indices
+
+ def _is_subtitle_row(self, row_element: Tag, row: ParsedRow, expected_width: int, real_cells: int) -> bool:
+ """
+ Проверяет, является ли строка подзаголовком для подтаблицы.
+ Подзаголовком считается строка с одной ячейкой, объединенной на всю ширину таблицы.
+
+ Args:
+ row_element (Tag): XML элемент строки
+ row (ParsedRow): Распарсенная строка
+ expected_width (int): Ожидаемая ширина таблицы
+ real_cells (int): Реальное количество ячеек в строке
+
+ Returns:
+ bool: True, если строка является подзаголовком
+ """
+ # Проверяем, содержит ли строка одну ячейку
+ if real_cells != 1:
+ return False
+
+ # Проверяем, что ячейка не пустая
+ cell_text = row.cells[0].strip()
+ if not cell_text:
+ return False
+
+ # Проверка 1: Непосредственная проверка gridSpan (объединенная ячейка)
+ grid_span = row_element.find('w:gridSpan')
+ if grid_span:
+ span_value = int(grid_span['w:val'])
+ # Ячейка объединена на всю ширину или близко к этому
+ if span_value >= expected_width - 1 and expected_width > 1:
+ return True
+
+ # Проверка 2: Если у нас есть только одна ячейка, а таблица шире,
+ # значит это, вероятно, подзаголовок
+ if expected_width > 1:
+ return True
+
+ return False
+
+ def _is_note_row(self, row_element: Tag, row: ParsedRow, expected_width: int, real_cells: int) -> bool:
+ """
+ Проверяет, является ли строка примечанием к таблице.
+ Примечанием считается последняя строка с одной ячейкой, объединенной на всю ширину таблицы.
+
+ Args:
+ row_element (Tag): XML элемент строки
+ row (ParsedRow): Распарсенная строка
+ expected_width (int): Ожидаемая ширина таблицы
+
+ Returns:
+ bool: True, если строка является примечанием
+ """
+ # Проверяем, содержит ли строка одну ячейку
+ if real_cells != 1:
+ return False
+
+ # Проверяем, есть ли в ячейке объединение на всю ширину
+ cell_element = row_element.find('w:tc')
+ if not cell_element:
+ return False
+
+ # Проверяем объединение ячеек
+ tc_pr = cell_element.find('w:tcPr')
+ if tc_pr:
+ grid_span = tc_pr.find('w:gridSpan')
+ # Если ячейка объединена на несколько столбцов и таблица многоколоночная
+ if grid_span and 'w:val' in grid_span.attrs and expected_width > 1:
+ span_value = int(grid_span['w:val'])
+ if span_value >= expected_width - 1:
+ cell_text = row.cells[0].strip()
+ # Если текст начинается с "Примечание" или звездочки, это почти наверняка примечание
+ if cell_text.startswith('Примечани') or cell_text.startswith('*'):
+ return True
+ # Или если это просто объединенная на всю ширину ячейка в конце таблицы
+ return True
+
+ # Проверяем текст - если строка явно выглядит как примечание
+ cell_text = row.cells[0].strip()
+ if cell_text.startswith('Примечани') or cell_text.startswith('*'):
+ return True
+
+ return False
+
+ def _split_into_subtables(self, rows_info: list[RowInfo]) -> list[ParsedSubtable]:
+ """
+ Разбивает строки таблицы на подтаблицы по подзаголовкам.
+
+ Args:
+ rows_info (list[RowInfo]): Список информации о строках таблицы
+
+ Returns:
+ list[ParsedSubtable]: Список подтаблиц
+ """
+ if not rows_info:
+ return []
+
+ subtables = []
+ current_subtable = ParsedSubtable()
+
+ # Удаляем строки, которые являются заголовками из основного набора строк
+ content_rows = [r for r in rows_info if not r.is_header and not r.is_note]
+
+ # Проходим по всем строкам и группируем их в подтаблицы
+ for row_info in content_rows:
+ if row_info.is_subtitle:
+ # Если у нас есть накопленные строки в текущей подтаблице, добавляем её
+ if current_subtable.rows:
+ subtables.append(current_subtable)
+
+ # Создаем новую подтаблицу с текущим подзаголовком
+ current_title = row_info.parsed_row.cells[0].strip()
+ current_subtable = ParsedSubtable(title=current_title)
+ else:
+ # Добавляем обычную строку в текущую подтаблицу
+ current_subtable.rows.append(row_info.parsed_row)
+
+ # Добавляем последнюю подтаблицу, если она содержит строки
+ if current_subtable.rows:
+ subtables.append(current_subtable)
+
+ # Если нет подтаблиц, создаем одну без названия
+ if not subtables and content_rows:
+ default_subtable = ParsedSubtable()
+ for row_info in content_rows:
+ if not row_info.is_subtitle:
+ default_subtable.rows.append(row_info.parsed_row)
+
+ if default_subtable.rows:
+ subtables.append(default_subtable)
+
+ return subtables
+
+ def _normalize_row_widths(self, headers: list[ParsedRow], subtables: list[ParsedSubtable], expected_width: int) -> None:
+ """
+ Нормализует ширину всех строк, чтобы она соответствовала ожидаемой ширине.
+
+ Args:
+ headers (list[ParsedRow]): Список строк заголовка
+ subtables (list[ParsedSubtable]): Список подтаблиц
+ expected_width (int): Ожидаемая ширина строки
+ """
+ # Если нет ожидаемой ширины, пытаемся определить её
+ if expected_width == 0:
+ # Определяем ширину на основе заголовков
+ if headers:
+ expected_width = len(headers[0].cells)
+ # Если заголовка нет, определяем ширину на основе первой строки первой подтаблицы
+ elif subtables and subtables[0].rows:
+ expected_width = max(len(row.cells) for row in subtables[0].rows)
+
+ # Если всё ещё нет ширины, нечего нормализовать
+ if expected_width == 0:
+ return
+
+ # Нормализуем строки заголовка
+ for header in headers:
+ self._normalize_single_row(header, expected_width)
+
+ # Нормализуем строки подтаблиц
+ for subtable in subtables:
+ for row in subtable.rows:
+ self._normalize_single_row(row, expected_width)
+
+ def _normalize_single_row(self, row: ParsedRow, expected_width: int) -> None:
+ """
+ Нормализует ширину одной строки.
+
+ Args:
+ row (ParsedRow): Строка для нормализации
+ expected_width (int): Ожидаемая ширина строки
+ """
+ current_width = len(row.cells)
+
+ # Добавляем пустые ячейки, если строка короче ожидаемой
+ if current_width < expected_width:
+ row.cells.extend([""] * (expected_width - current_width))
+ # Обрезаем лишние ячейки, если строка длиннее ожидаемой
+ elif current_width > expected_width:
+ row.cells = row.cells[:expected_width]
+
+ def _parse_row(self, row_element: Tag) -> tuple[ParsedRow, int]:
+ """
+ Парсит строку таблицы, обрабатывая горизонтальные объединения ячеек.
+
+ Args:
+ row_element (Tag): XML элемент строки
+
+ Returns:
+ tuple[ParsedRow, int]: Кортеж с данными строки и реальным количеством ячеек
+ """
+ cells_text = []
+ cell_elements = row_element.find_all('w:tc')
+ row_index = 0 # Для индексации строки
+ real_cells = len(cell_elements)
+
+ # Извлекаем текст из всех ячеек строки, обрабатывая горизонтальные объединения
+ for cell_element in cell_elements:
+ # Получаем текст ячейки
+ cell_text = self._extract_cell_text(cell_element)
+
+ # Извлекаем информацию о горизонтальном объединении
+ h_span = 1
+ tc_pr = cell_element.find('w:tcPr')
+ if tc_pr:
+ grid_span = tc_pr.find('w:gridSpan')
+ if grid_span and 'w:val' in grid_span.attrs:
+ h_span = int(grid_span['w:val'])
+
+ # Дублируем ячейку нужное количество раз для учета горизонтального объединения
+ for _ in range(h_span):
+ cells_text.append(cell_text)
+
+ # Собираем информацию о стиле строки
+ row_style = self._extract_row_style(row_element)
+
+ return ParsedRow(
+ index=row_index,
+ cells=cells_text,
+ style=row_style
+ ), real_cells
+
+ def _extract_cell_text(self, cell_element: Tag) -> str:
+ """
+ Извлекает текст из ячейки.
+
+ Args:
+ cell_element (Tag): XML элемент ячейки
+
+ Returns:
+ str: Текст ячейки
+ """
+ # Извлекаем текст из всех параграфов в ячейке
+ paragraphs = cell_element.find_all('w:p')
+
+ paragraph_texts = []
+ for p in paragraphs:
+ p_text = ""
+ for run in p.find_all('w:r'):
+ for t in run.find_all('w:t'):
+ p_text += t.get_text()
+ paragraph_texts.append(p_text)
+
+ # Объединяем текст параграфов с переносом строки
+ return "\n".join(paragraph_texts)
+
+ def _extract_row_style(self, row_element: Tag) -> TextStyle:
+ """
+ Извлекает информацию о стиле строки.
+
+ Args:
+ row_element (Tag): XML элемент строки
+
+ Returns:
+ TextStyle: Стиль строки
+ """
+ style = TextStyle()
+
+ # Счетчики для определения стиля всей строки
+ bold_runs = 0
+ italic_runs = 0
+ total_runs = 0
+ center_aligned_paragraphs = 0
+ total_paragraphs = 0
+
+ # Анализируем стиль всех ячеек
+ for cell in row_element.find_all('w:tc'):
+ for p in cell.find_all('w:p'):
+ total_paragraphs += 1
+
+ # Проверяем выравнивание параграфа
+ p_pr = p.find('w:pPr')
+ if p_pr:
+ jc = p_pr.find('w:jc')
+ if jc and 'w:val' in jc.attrs and jc['w:val'] == 'center':
+ center_aligned_paragraphs += 1
+
+ # Проверяем стиль текста в параграфе
+ for run in p.find_all('w:r'):
+ total_runs += 1
+ r_pr = run.find('w:rPr')
+ if r_pr:
+ if r_pr.find('w:b'):
+ bold_runs += 1
+ if r_pr.find('w:i'):
+ italic_runs += 1
+
+ # Устанавливаем стиль на основе анализа
+ if total_runs > 0:
+ if bold_runs == total_runs:
+ style.fully_bold = True
+ elif bold_runs > 0:
+ style.partly_bold = True
+
+ if italic_runs == total_runs:
+ style.fully_italic = True
+ elif italic_runs > 0:
+ style.partly_italic = True
+
+ # Устанавливаем выравнивание
+ if total_paragraphs > 0 and center_aligned_paragraphs >= total_paragraphs * 0.8:
+ style.alignment = 'center'
+
+ return style
diff --git a/lib/parser/ntr_fileparser/parsers/specific_parsers/xml_parser.py b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..6527bc3a4fea01906ad287661e22472bf064c1dc
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/specific_parsers/xml_parser.py
@@ -0,0 +1,226 @@
+"""
+Модуль с парсером для XML документов.
+"""
+
+import logging
+import os
+import re
+from typing import Any, BinaryIO
+
+from bs4 import BeautifulSoup
+
+from ...data_classes import ParsedDocument
+from ..abstract_parser import AbstractParser
+from ..file_types import FileType
+from .xml.formula_parser import XMLFormulaParser
+from .xml.image_parser import XMLImageParser
+from .xml.meta_parser import XMLMetaParser
+from .xml.paragraph_parser import XMLParagraphParser
+from .xml.table_parser import XMLTableParser
+
+logger = logging.getLogger(__name__)
+
+
+class XMLParser(AbstractParser):
+ """
+ Парсер для XML документов.
+
+ Поддерживает извлечение текста, таблиц, формул и других элементов
+ из XML файлов, используя BeautifulSoup.
+ """
+
+ def __init__(
+ self,
+ style_cache: dict[str, Any] | None = None,
+ numbering_cache: dict[str, Any] | None = None,
+ relationships_cache: dict[str, Any] | None = None,
+ ):
+ """
+ Инициализирует XML парсер и его компоненты.
+
+ Args:
+ style_cache (dict[str, Any] | None): Кэш стилей для передачи парсеру параграфов
+ numbering_cache (dict[str, Any] | None): Кэш нумерации для передачи парсеру параграфов
+ relationships_cache (dict[str, Any] | None): Кэш связей для обработки референсов
+ """
+ super().__init__(FileType.XML)
+ self.table_parser = XMLTableParser()
+ self.paragraph_parser = XMLParagraphParser(style_cache, numbering_cache, relationships_cache)
+ self.image_parser = XMLImageParser()
+ self.formula_parser = XMLFormulaParser()
+ self.meta_parser = XMLMetaParser()
+
+ def _detect_encoding(self, content: bytes) -> str:
+ """
+ Определяет кодировку из XML заголовка или возвращает cp866.
+
+ Args:
+ content (bytes): Содержимое XML файла.
+
+ Returns:
+ str: Определенная кодировка или cp866 по умолчанию.
+ """
+ try:
+ # Пытаемся прочитать первые строки как UTF-8 для поиска объявления XML
+ header = content[:1000].decode('utf-8', errors='ignore')
+ if ' ParsedDocument:
+ """
+ Парсит XML документ по пути к файлу и возвращает его структурное представление.
+
+ Args:
+ file_path (str): Путь к XML файлу для парсинга.
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если файл не может быть прочитан или распарсен.
+ """
+ logger.debug(f"Parsing XML file: {file_path}")
+
+ if not os.path.exists(file_path):
+ raise ValueError(f"File not found: {file_path}")
+
+ with open(file_path, 'rb') as f:
+ content = f.read()
+
+ # Извлекаем имя файла из пути
+ filename = os.path.basename(file_path)
+
+ return self._parse_content(content, filename, file_path)
+
+ def parse(
+ self,
+ file: BinaryIO,
+ file_type: FileType | str | None = None,
+ ) -> ParsedDocument:
+ """
+ Парсит XML документ из объекта файла и возвращает его структурное представление.
+
+ Args:
+ file (BinaryIO): Объект файла для парсинга.
+ file_type: Тип файла, если известен.
+ Может быть объектом FileType или строкой с расширением (".xml").
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если файл не может быть прочитан или распарсен.
+ """
+ logger.debug("Parsing XML from file object")
+
+ if file_type and isinstance(file_type, FileType) and file_type != FileType.XML:
+ logger.warning(
+ f"Provided file_type {file_type} doesn't match parser type {FileType.XML}"
+ )
+
+ # Читаем содержимое файла
+ content = file.read()
+
+ return self._parse_content(content, "unknown.xml", None)
+
+ def _parse_content(
+ self,
+ content: bytes,
+ filename: str,
+ filepath: str | None,
+ ) -> ParsedDocument:
+ """
+ Внутренний метод для парсинга содержимого XML файла.
+
+ Args:
+ content (bytes): Содержимое XML файла.
+ filename (str): Имя файла для документа.
+ filepath (str | None): Путь к файлу (или None, если из объекта).
+
+ Returns:
+ ParsedDocument: Структурное представление документа.
+
+ Raises:
+ ValueError: Если содержимое не может быть распарсено.
+ """
+ # Определение кодировки из XML заголовка
+ encoding = self._detect_encoding(content)
+ logger.debug(f"Detected encoding: {encoding}")
+
+ try:
+ xml_text = content.decode(encoding)
+ except UnicodeDecodeError as e:
+ logger.error(f"Failed to decode XML with {encoding} encoding: {e}")
+ raise ValueError(f"Cannot decode XML content with {encoding} encoding")
+
+ # Создание BeautifulSoup один раз
+ try:
+ soup = BeautifulSoup(xml_text, features='xml')
+ logger.debug("Created BeautifulSoup object")
+ except Exception as e:
+ logger.error(f"Failed to parse XML: {e}")
+ raise ValueError("Cannot parse XML content")
+
+ # Создание базового документа
+ doc = ParsedDocument(name=filename, type="XML")
+
+ # Извлечение метаданных
+ doc.meta = self.meta_parser.parse(soup, filepath)
+ logger.debug("Parsed metadata")
+
+ # Последовательный вызов парсеров
+ try:
+ # Вызываем парсеры, которые не модифицируют soup
+ doc.tables.extend(self.table_parser.parse(soup))
+ logger.debug(f"Parsed {len(doc.tables)} tables")
+
+ doc.images.extend(self.image_parser.parse(soup))
+ logger.debug(f"Parsed {len(doc.images)} images")
+
+ doc.formulas.extend(self.formula_parser.parse(soup))
+ logger.debug(f"Parsed {len(doc.formulas)} formulas")
+
+ # Вызываем парсер параграфов последним, т.к. он модифицирует soup
+ # (удаляет таблицы, изображения и др. элементы)
+ doc.paragraphs.extend(self.paragraph_parser.parse(soup))
+ logger.debug(f"Parsed {len(doc.paragraphs)} paragraphs")
+
+ # Связываем элементы на основе полнотекстового совпадения
+ self._link_elements(doc)
+ logger.debug("Linked elements based on text matching")
+ except Exception as e:
+ logger.error(f"Error during parsing components: {e}")
+ raise ValueError("Error parsing document components")
+
+ return doc
+
+ def _link_elements(self, doc: ParsedDocument) -> None:
+ """
+ Связывает таблицы, изображения и формулы с соответствующими параграфами
+ на основе полнотекстового совпадения.
+
+ Args:
+ doc (ParsedDocument): Документ для обработки.
+ """
+ # Индексируем параграфы документа
+ for i, paragraph in enumerate(doc.paragraphs):
+ paragraph.index_in_document = i
+
+ for i, table in enumerate(doc.tables):
+ table.index_in_document = i
+
+ last_index = 0
+
+ for table in doc.tables:
+ for i in range(last_index, len(doc.paragraphs)):
+ paragraph = doc.paragraphs[i]
+ if table.title == paragraph.text:
+ table.title_index_in_paragraphs = i
+ paragraph.title_of_table = table.index_in_document
+ last_index = i
+ break
diff --git a/lib/parser/ntr_fileparser/parsers/universal_parser.py b/lib/parser/ntr_fileparser/parsers/universal_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..108b14133e0e9abeeae6a3258aa1708c40943eaa
--- /dev/null
+++ b/lib/parser/ntr_fileparser/parsers/universal_parser.py
@@ -0,0 +1,149 @@
+"""
+Модуль с универсальным парсером, объединяющим все специфичные парсеры.
+"""
+
+import logging
+import os
+from typing import BinaryIO
+
+from ..data_classes import ParsedDocument
+from .abstract_parser import AbstractParser
+from .file_types import FileType
+from .parser_factory import ParserFactory
+from .specific_parsers import (
+ DocParser,
+ DocxParser,
+ EmailParser,
+ HTMLParser,
+ MarkdownParser,
+ PDFParser,
+ XMLParser,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class UniversalParser:
+ """
+ Универсальный парсер, объединяющий все специфичные парсеры.
+
+ Использует фабрику парсеров для выбора подходящего парсера
+ на основе типа файла.
+ """
+
+ def __init__(self):
+ """
+ Инициализирует универсальный парсер и регистрирует все доступные парсеры.
+ """
+ self.factory = ParserFactory()
+
+ # Регистрируем все доступные парсеры
+ self.register_parsers(
+ [
+ XMLParser(), # Реализованный парсер
+ PDFParser(), # Нереализованный парсер
+ DocParser(), # Нереализованный парсер
+ DocxParser(), # Реализованный парсер
+ EmailParser(), # Нереализованный парсер
+ MarkdownParser(), # Нереализованный парсер
+ HTMLParser(), # Нереализованный парсер
+ ]
+ )
+
+ def register_parser(self, parser: AbstractParser) -> None:
+ """
+ Регистрирует парсер в фабрике.
+
+ Args:
+ parser (AbstractParser): Парсер для регистрации.
+ """
+ self.factory.register_parser(parser)
+
+ def register_parsers(self, parsers: list[AbstractParser]) -> None:
+ """
+ Регистрирует несколько парсеров в фабрике.
+
+ Args:
+ parsers (list[AbstractParser]): Список парсеров для регистрации.
+ """
+ for parser in parsers:
+ self.register_parser(parser)
+
+ def parse_by_path(self, file_path: str) -> ParsedDocument | None:
+ """
+ Парсит документ по пути к файлу, используя подходящий парсер.
+
+ Args:
+ file_path (str): Путь к файлу для парсинга.
+
+ Returns:
+ ParsedDocument | None: Структурное представление документа или None,
+ если подходящий парсер не найден.
+
+ Raises:
+ ValueError: Если файл не существует или не может быть прочитан.
+ """
+ if not os.path.exists(file_path):
+ raise ValueError(f"Файл не найден: {file_path}")
+
+ # Находим подходящий парсер
+ parser = self.factory.get_parser(file_path)
+ if not parser:
+ logger.warning(f"Не найден подходящий парсер для файла: {file_path}")
+ return None
+
+ # Парсим документ
+ try:
+ return parser.parse_by_path(file_path)
+ except Exception as e:
+ logger.error(f"Ошибка при парсинге файла {file_path}: {e}")
+ raise
+
+ def parse(
+ self, file: BinaryIO, file_type: FileType | str | None = None
+ ) -> ParsedDocument | None:
+ """
+ Парсит документ из объекта файла, используя подходящий парсер.
+
+ Args:
+ file (BinaryIO): Объект файла для парсинга.
+ file_type: Тип файла, может быть объектом FileType или строкой с расширением.
+ Например: FileType.XML или ".xml"
+
+ Returns:
+ ParsedDocument | None: Структурное представление документа или None,
+ если подходящий парсер не найден.
+
+ Raises:
+ ValueError: Если файл не может быть прочитан или распарсен.
+ """
+ # Преобразуем строковое расширение в FileType, если нужно
+ ft = None
+ if isinstance(file_type, str):
+ try:
+ ft = FileType.from_extension(file_type)
+ except ValueError:
+ logger.warning(f"Неизвестное расширение файла: {file_type}")
+ return None
+ else:
+ ft = file_type
+
+ if ft is None:
+ logger.warning("Тип файла не указан при парсинге из объекта файла")
+ return None
+
+ # Получаем парсер для указанного типа файла
+ parsers = [p for p in self.factory.parsers if p.supports_file(ft)]
+ if not parsers:
+ logger.warning(f"Не найден подходящий парсер для типа файла: {ft}")
+ return None
+
+ # Используем первый подходящий парсер
+ parser = parsers[0]
+
+ # Парсим документ
+ try:
+ return parser.parse(file, ft)
+ except Exception as e:
+ logger.error(f"Ошибка при парсинге файла: {e}")
+ raise
diff --git a/lib/parser/ntr_fileparser_diagram.puml b/lib/parser/ntr_fileparser_diagram.puml
new file mode 100644
index 0000000000000000000000000000000000000000..780cd96eb7ea32dede1010f2a7850c4f4c9c2d8a
--- /dev/null
+++ b/lib/parser/ntr_fileparser_diagram.puml
@@ -0,0 +1,221 @@
+@startuml NTR_FileParser
+
+package "ntr_fileparser" {
+ package "data_classes" {
+ abstract class ParsedStructure {
+ +{abstract} apply(func: Callable[[str], str])
+ +{abstract} to_dict()
+ +{abstract} to_string()
+ }
+
+ class ParsedDocument {
+ +name: str
+ +type: str
+ +meta: ParsedMeta
+ +paragraphs: list[ParsedTextBlock]
+ +tables: list[ParsedTable]
+ +images: list[ParsedImage]
+ +formulas: list[ParsedFormula]
+ }
+
+ class ParsedMeta {
+ +title: str
+ +author: str
+ +creation_date: str
+ }
+
+ class ParsedTextBlock {
+ +text: str
+ +style: TextStyle
+ }
+
+ enum TextStyle {
+ NORMAL
+ BOLD
+ ITALIC
+ UNDERLINE
+ HEADING1
+ HEADING2
+ HEADING3
+ }
+
+ class ParsedTable {
+ +headers: list[str]
+ +rows: list[ParsedRow]
+ +subtables: list[ParsedSubtable]
+ +tag: TableTag
+ }
+
+ class ParsedRow {
+ +cells: list[str]
+ }
+
+ class ParsedSubtable {
+ +table: ParsedTable
+ }
+
+ enum TableTag {
+ UNKNOWN
+ DATA
+ METADATA
+ }
+
+ class ParsedImage #lightgrey {
+ +path: str
+ +alt_text: str
+ .. Примечание ..
+ В текущей реализации не используется
+ }
+
+ class ParsedFormula #lightgrey {
+ +latex: str
+ .. Примечание ..
+ В текущей реализации не используется
+ }
+
+ ParsedStructure <|-- ParsedDocument
+ ParsedStructure <|-- ParsedTextBlock
+ ParsedStructure <|-- ParsedTable
+ ParsedStructure <|-- ParsedRow
+ ParsedStructure <|-- ParsedSubtable
+ ParsedStructure <|-- ParsedImage
+ ParsedStructure <|-- ParsedFormula
+ ParsedStructure <|-- ParsedMeta
+
+ ParsedDocument o-- ParsedMeta
+ ParsedDocument o-- "*" ParsedTextBlock
+ ParsedDocument o-- "*" ParsedTable
+ ParsedDocument o-- "*" ParsedImage
+ ParsedDocument o-- "*" ParsedFormula
+ ParsedTable o-- "*" ParsedRow
+ ParsedTable o-- "*" ParsedSubtable
+ ParsedTable -- TableTag
+ ParsedTextBlock -- TextStyle
+ }
+
+ package "parsers" {
+ abstract class AbstractParser {
+ +file_types: list
+ +{abstract} parse()
+ +{abstract} parse_by_path()
+ +supports_file()
+ +_supported_extension()
+ }
+
+ class ParserFactory {
+ +parsers: list[AbstractParser]
+ +register_parser()
+ +get_parser()
+ }
+
+ class UniversalParser {
+ +factory: ParserFactory
+ +parse()
+ +parse_by_path()
+ }
+
+ enum FileType {
+ XML
+ DOCX
+ DOC
+ PDF
+ HTML
+ MD
+ EML
+ +from_extension()
+ +get_supported_extensions()
+ }
+
+ package "specific_parsers" {
+ package "xml" {
+ class XMLParagraphParser {
+ +parse()
+ }
+
+ class XMLTableParser {
+ +parse()
+ }
+
+ class XMLMetaParser {
+ +parse()
+ +_extract_info_value()
+ +_extract_info_recurse()
+ }
+
+ class XMLImageParser #lightgrey {
+ +parse()
+ .. Примечание ..
+ В текущей реализации не используется
+ }
+
+ class XMLFormulaParser #lightgrey {
+ +parse()
+ .. Примечание ..
+ В текущей реализации не используется
+ }
+ }
+
+ package "docx" {
+ class CorePropertiesParser {
+ +parse()
+ }
+
+ class MetadataParser {
+ +parse()
+ }
+
+ class NumberingParser {
+ +parse()
+ }
+
+ class RelationshipsParser {
+ +parse()
+ }
+
+ class StylesParser {
+ +parse()
+ }
+ }
+
+ class DocParser {
+ }
+
+ class DocxParser {
+ }
+
+ class PDFParser {
+ }
+
+ class XMLParser {
+ }
+
+ class HTMLParser {
+ }
+
+ class MarkdownParser {
+ }
+
+ class EmailParser {
+ }
+
+ XMLParser -- xml
+ DocxParser -- docx
+ }
+
+ AbstractParser <|-- DocParser
+ AbstractParser <|-- DocxParser
+ AbstractParser <|-- PDFParser
+ AbstractParser <|-- XMLParser
+ AbstractParser <|-- HTMLParser
+ AbstractParser <|-- MarkdownParser
+ AbstractParser <|-- EmailParser
+
+ AbstractParser -- FileType
+ ParserFactory o-- "*" AbstractParser
+ UniversalParser --> ParserFactory
+ }
+
+ data_classes <.. parsers : использует
+}
+
+@enduml
\ No newline at end of file
diff --git a/lib/parser/pyproject.toml b/lib/parser/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..8964aaf73c849aad9eab184e30ce678748ad705e
--- /dev/null
+++ b/lib/parser/pyproject.toml
@@ -0,0 +1,16 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = ["setuptools>=61"]
+
+[project]
+name = "ntr_fileparser"
+version = "0.2.0"
+dependencies = [
+ "beautifulsoup4>=4.11.1",
+ "lxml>=4.9.1",
+ "typing-extensions>=4.4.0",
+ "PyMuPDF>=1.21.0",
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
diff --git a/lib/parser/scripts/test_docx.py b/lib/parser/scripts/test_docx.py
new file mode 100644
index 0000000000000000000000000000000000000000..ccffc09032183fcb070ad5146ca3895f5b1e7d08
--- /dev/null
+++ b/lib/parser/scripts/test_docx.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+"""
+Скрипт для тестирования парсера DOCX документов
+и сохранения результатов в JSON формате.
+
+Положите ваш файл test.docx в директорию test_input и запустите скрипт.
+Результат будет сохранен в файл test_output/test.json
+"""
+
+from datetime import datetime
+import json
+import logging
+import sys
+from pathlib import Path
+
+# Добавляем родительскую директорию в sys.path, чтобы импорты работали корректно
+# при запуске скрипта из верхнего уровня
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+from parser import UniversalParser # type: ignore
+
+logging.basicConfig(level=logging.INFO, handlers=[logging.StreamHandler()])
+
+
+def main():
+ """
+ Основная функция скрипта.
+
+ 1. Считывает файл test_input/test.docx
+ 2. Парсит его с помощью UniversalParser
+ 3. Сохраняет результат в test_output/test.json
+ """
+ # Получаем абсолютные пути к файлам относительно корневой директории проекта
+ project_root = Path(__file__).parent.parent
+ input_file = project_root / "test_input" / "test.docx"
+ output_file = project_root / "test_output" / "test.json"
+
+ # Проверяем существование входного файла
+ if not input_file.exists():
+ print(
+ f"Ошибка: Файл {input_file} не найден. Пожалуйста, поместите файл test.docx в директорию test_input."
+ )
+ return 1
+
+ # Создаем директорию для выходных файлов, если она не существует
+ output_file.parent.mkdir(parents=True, exist_ok=True)
+
+ print(f"Парсинг файла: {input_file}")
+
+ # Создаем экземпляр универсального парсера
+ parser = UniversalParser()
+
+ # Парсим документ
+ parsed_document = parser.parse_by_path(str(input_file))
+
+ if parsed_document is None:
+ print("Ошибка: Не удалось распарсить документ.")
+ return 1
+
+ # Преобразуем документ в словарь
+ document_dict = parsed_document.to_dict()
+
+ # Сохраняем результат в JSON
+ with open(output_file, "w", encoding="utf-8") as f:
+ json.dump(document_dict, f, ensure_ascii=False, indent=2)
+
+ print(f"Результат сохранен в файле: {output_file}")
+ return 0
+
+
+if __name__ == "__main__":
+ time_start = datetime.now()
+ for i in range(3):
+ main()
+ time_end = datetime.now()
+ print(f"Время выполнения: {(time_end - time_start).total_seconds()} секунд")
diff --git a/lib/parser/scripts/test_pdf.py b/lib/parser/scripts/test_pdf.py
new file mode 100644
index 0000000000000000000000000000000000000000..7639d810da9b9c8b00e5cba578aebc1e95defadd
--- /dev/null
+++ b/lib/parser/scripts/test_pdf.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+"""
+Скрипт для тестирования парсера PDF документов
+и сохранения результатов в JSON формате.
+
+Положите ваш файл test.pdf в директорию test_input и запустите скрипт.
+Результат будет сохранен в файл test_output/test_pdf.json
+"""
+
+import json
+import logging
+import sys
+from datetime import datetime
+from pathlib import Path
+
+# Добавляем родительскую директорию в sys.path, чтобы импорты работали корректно
+# при запуске скрипта из верхнего уровня
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+from parser import UniversalParser # type: ignore
+
+logging.basicConfig(level=logging.INFO, handlers=[logging.StreamHandler()])
+
+
+def main():
+ """
+ Основная функция скрипта.
+
+ 1. Считывает файл test_input/test.pdf
+ 2. Парсит его с помощью UniversalParser
+ 3. Сохраняет результат в test_output/test_pdf.json
+ """
+ # Получаем абсолютные пути к файлам относительно корневой директории проекта
+ project_root = Path(__file__).parent.parent
+ input_file = project_root / "test_input" / "test.pdf"
+ output_file = project_root / "test_output" / "test_pdf.json"
+
+ # Проверяем существование входного файла
+ if not input_file.exists():
+ print(
+ f"Ошибка: Файл {input_file} не найден. Пожалуйста, поместите файл test.pdf в директорию test_input."
+ )
+ return 1
+
+ # Создаем директорию для выходных файлов, если она не существует
+ output_file.parent.mkdir(parents=True, exist_ok=True)
+
+ print(f"Парсинг PDF-файла: {input_file}")
+
+ # Создаем экземпляр универсального парсера
+ parser = UniversalParser()
+
+ # Парсим документ
+ parsed_document = parser.parse_by_path(str(input_file))
+
+ if parsed_document is None:
+ print("Ошибка: Не удалось распарсить PDF-документ.")
+ return 1
+
+ # Преобразуем документ в словарь
+ document_dict = parsed_document.to_dict()
+
+ # Сохраняем результат в JSON
+ with open(output_file, "w", encoding="utf-8") as f:
+ json.dump(document_dict, f, ensure_ascii=False, indent=2)
+
+ print(f"Результат сохранен в файле: {output_file}")
+ return 0
+
+
+if __name__ == "__main__":
+ time_start = datetime.now()
+ for i in range(1):
+ main()
+ time_end = datetime.now()
+ print(f"Время выполнения: {(time_end - time_start).total_seconds()} секунд")
\ No newline at end of file
diff --git a/main.py b/main.py
index 841ed46f3a8781f0b3aaa5753501698c061b74e4..ea05c03ee3d635e0771e186f591ba99be55d6417 100644
--- a/main.py
+++ b/main.py
@@ -1,27 +1,28 @@
+import logging
+import os
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated
+
import dotenv
import uvicorn
-import logging
-import os
-from fastapi import FastAPI, Depends
+from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
-
-from common.common import configure_logging
-from common.configuration import Configuration
-# from main_before import config
+from transformers import AutoModel, AutoTokenizer
# from routes.acronym import router as acronym_router
from common import dependencies as DI
+from common.common import configure_logging
+from common.configuration import Configuration
from routes.dataset import router as dataset_router
from routes.document import router as document_router
-from routes.acronym import router as acronym_router
+from routes.entity import router as entity_router
from routes.llm import router as llm_router
from routes.llm_config import router as llm_config_router
from routes.llm_prompt import router as llm_prompt_router
-from common.common import configure_logging
-from transformers import AutoTokenizer, AutoModel
+
+# from main_before import config
+
# Загружаем переменные из .env
dotenv.load_dotenv()
@@ -67,11 +68,11 @@ app.add_middleware(
app.include_router(llm_router)
# app.include_router(log_router)
# app.include_router(feedback_router)
-app.include_router(acronym_router)
app.include_router(dataset_router)
app.include_router(document_router)
app.include_router(llm_config_router)
app.include_router(llm_prompt_router)
+app.include_router(entity_router)
if __name__ == "__main__":
uvicorn.run(
diff --git a/routes/dataset.py b/routes/dataset.py
index 64d7d3c0f59cf644e60efa8c0bed7849a16146e7..f374fccd398a9e0f295c5cb1de7e9f5ffe57fe7a 100644
--- a/routes/dataset.py
+++ b/routes/dataset.py
@@ -1,12 +1,13 @@
import logging
from typing import Annotated
-from fastapi import APIRouter, BackgroundTasks, HTTPException, Response, UploadFile, Depends
+from fastapi import (APIRouter, BackgroundTasks, Depends, HTTPException,
+ Response, UploadFile)
+import common.dependencies as DI
from components.services.dataset import DatasetService
from schemas.dataset import (Dataset, DatasetExpanded, DatasetProcessing,
SortQuery, SortQueryList)
-import common.dependencies as DI
router = APIRouter(prefix='/datasets')
logger = logging.getLogger(__name__)
@@ -48,7 +49,7 @@ def try_create_default_dataset(dataset_service: DatasetService):
else:
dataset_service.create_dataset_from_directory(
is_default=True,
- directory_with_xmls=dataset_service.config.db_config.files.xmls_path_default,
+ directory_with_documents=dataset_service.config.db_config.files.xmls_path_default,
directory_with_ready_dataset=dataset_service.config.db_config.files.start_path,
)
diff --git a/routes/entity.py b/routes/entity.py
new file mode 100644
index 0000000000000000000000000000000000000000..3fb3ef9073a83c3b264c5d5dcd1b140e44cc83cd
--- /dev/null
+++ b/routes/entity.py
@@ -0,0 +1,264 @@
+from typing import Annotated
+
+import numpy as np
+from fastapi import APIRouter, Depends, HTTPException
+from sqlalchemy.orm import Session
+
+import common.dependencies as DI
+from components.dbo.chunk_repository import ChunkRepository
+from components.services.entity import EntityService
+from schemas.entity import (EntityNeighborsRequest, EntityNeighborsResponse,
+ EntitySearchRequest, EntitySearchResponse,
+ EntitySearchWithTextRequest,
+ EntitySearchWithTextResponse, EntityTextRequest,
+ EntityTextResponse)
+
+router = APIRouter(prefix="/entity", tags=["entity"])
+
+
+@router.post("/search", response_model=EntitySearchResponse)
+async def search_entities(
+ request: EntitySearchRequest,
+ entity_service: EntityService = Depends(DI.get_entity_service),
+) -> EntitySearchResponse:
+ """
+ Поиск похожих сущностей по векторному сходству (только ID).
+
+ Args:
+ request: Параметры поиска
+ entity_service: Сервис для работы с сущностями
+
+ Returns:
+ Результаты поиска (ID и оценки), отсортированные по убыванию сходства
+ """
+ try:
+ _, scores, ids = entity_service.search_similar(
+ request.query,
+ request.dataset_id,
+ )
+
+ # Проверяем, что scores и ids - корректные numpy массивы
+ if not isinstance(scores, np.ndarray):
+ scores = np.array(scores)
+ if not isinstance(ids, np.ndarray):
+ ids = np.array(ids)
+
+ # Сортируем результаты по убыванию оценок
+ # Проверим, что массивы не пустые
+ if len(scores) > 0:
+ # Преобразуем индексы в список, чтобы избежать проблем с индексацией
+ sorted_indices = scores.argsort()[::-1].tolist()
+ sorted_scores = [float(scores[i]) for i in sorted_indices]
+ # Преобразуем все ID в строки
+ sorted_ids = [str(ids[i]) for i in sorted_indices]
+ else:
+ sorted_scores = []
+ sorted_ids = []
+
+ return EntitySearchResponse(
+ scores=sorted_scores,
+ entity_ids=sorted_ids,
+ )
+ except Exception as e:
+ raise HTTPException(
+ status_code=500,
+ detail=f"Error during entity search: {str(e)}"
+ )
+
+
+@router.post("/search/with_text", response_model=EntitySearchWithTextResponse)
+async def search_entities_with_text(
+ request: EntitySearchWithTextRequest,
+ entity_service: EntityService = Depends(DI.get_entity_service),
+) -> EntitySearchWithTextResponse:
+ """
+ Поиск похожих сущностей по векторному сходству с возвратом текстов.
+
+ Args:
+ request: Параметры поиска
+ entity_service: Сервис для работы с сущностями
+
+ Returns:
+ Результаты поиска с текстами чанков, отсортированные по убыванию сходства
+ """
+ try:
+ # Получаем результаты поиска
+ _, scores, entity_ids = entity_service.search_similar(
+ request.query,
+ request.dataset_id
+ )
+
+ # Проверяем, что scores и entity_ids - корректные numpy массивы
+ if not isinstance(scores, np.ndarray):
+ scores = np.array(scores)
+ if not isinstance(entity_ids, np.ndarray):
+ entity_ids = np.array(entity_ids)
+
+ # Сортируем результаты по убыванию оценок
+ # Проверим, что массивы не пустые
+ if len(scores) > 0:
+ # Преобразуем индексы в список, чтобы избежать проблем с индексацией
+ sorted_indices = scores.argsort()[::-1].tolist()
+ sorted_scores = [float(scores[i]) for i in sorted_indices]
+ sorted_ids = [str(entity_ids[i]) for i in sorted_indices] # Преобразуем в строки
+
+ # Получаем тексты чанков
+ chunks = entity_service.chunk_repository.get_chunks_by_ids(sorted_ids)
+
+ # Формируем ответ
+ return EntitySearchWithTextResponse(
+ chunks=[
+ {
+ "id": str(chunk.id), # Преобразуем UUID в строку
+ "text": chunk.text,
+ "score": score
+ }
+ for chunk, score in zip(chunks, sorted_scores)
+ ]
+ )
+ else:
+ return EntitySearchWithTextResponse(chunks=[])
+
+ except Exception as e:
+ raise HTTPException(
+ status_code=500,
+ detail=f"Error during entity search with text: {str(e)}"
+ )
+
+
+@router.post("/text", response_model=EntityTextResponse)
+async def build_entity_text(
+ request: EntityTextRequest,
+ entity_service: EntityService = Depends(DI.get_entity_service),
+) -> EntityTextResponse:
+ """
+ Сборка текста из сущностей.
+
+ Args:
+ request: Параметры сборки текста
+ entity_service: Сервис для работы с сущностями
+
+ Returns:
+ Собранный текст
+ """
+ try:
+ # Получаем объекты LinkerEntity по ID
+ entities = entity_service.chunk_repository.get_chunks_by_ids(request.entities)
+
+ if not entities:
+ raise HTTPException(
+ status_code=404,
+ detail="No entities found with provided IDs"
+ )
+
+ # Собираем текст
+ text = entity_service.build_text(
+ entities=entities,
+ chunk_scores=request.chunk_scores,
+ include_tables=request.include_tables,
+ max_documents=request.max_documents,
+ )
+
+ return EntityTextResponse(text=text)
+ except Exception as e:
+ raise HTTPException(
+ status_code=500,
+ detail=f"Error building entity text: {str(e)}"
+ )
+
+
+@router.post("/neighbors", response_model=EntityNeighborsResponse)
+async def get_neighboring_chunks(
+ request: EntityNeighborsRequest,
+ entity_service: EntityService = Depends(DI.get_entity_service),
+) -> EntityNeighborsResponse:
+ """
+ Получение соседних чанков для заданных сущностей.
+
+ Args:
+ request: Параметры запроса соседей
+ entity_service: Сервис для работы с сущностями
+
+ Returns:
+ Список сущностей с соседями
+ """
+ try:
+ # Получаем объекты LinkerEntity по ID
+ entities = entity_service.chunk_repository.get_chunks_by_ids(request.entities)
+
+ if not entities:
+ raise HTTPException(
+ status_code=404,
+ detail="No entities found with provided IDs"
+ )
+
+ # Получаем соседние чанки
+ entities_with_neighbors = entity_service.add_neighboring_chunks(
+ entities,
+ max_distance=request.max_distance,
+ )
+
+ # Преобразуем LinkerEntity в строки
+ return EntityNeighborsResponse(
+ entities=[str(entity.id) for entity in entities_with_neighbors]
+ )
+ except Exception as e:
+ raise HTTPException(
+ status_code=500,
+ detail=f"Error getting neighboring chunks: {str(e)}"
+ )
+
+
+@router.get("/info/{dataset_id}")
+async def get_entity_info(
+ dataset_id: int,
+ db: Annotated[Session, Depends(DI.get_db)],
+) -> dict:
+ """
+ Получить информацию о сущностях в датасете.
+
+ Args:
+ dataset_id: ID датасета
+ db: Сессия базы данных
+ config: Конфигурация приложения
+
+ Returns:
+ dict: Информация о сущностях
+ """
+ chunk_repository = ChunkRepository(db)
+ entities, embeddings = chunk_repository.get_searching_entities(dataset_id)
+
+ if not entities:
+ raise HTTPException(status_code=404, detail=f"No entities found for dataset {dataset_id}")
+
+ # Собираем статистику
+ stats = {
+ "total_entities": len(entities),
+ "entities_with_embeddings": len([e for e in embeddings if e is not None]),
+ "embedding_shapes": [e.shape if e is not None else None for e in embeddings],
+ "unique_embedding_shapes": set(str(e.shape) if e is not None else None for e in embeddings),
+ "entity_types": set(e.type for e in entities),
+ "entities_per_type": {
+ t: len([e for e in entities if e.type == t])
+ for t in set(e.type for e in entities)
+ }
+ }
+
+ # Примеры сущностей
+ examples = [
+ {
+ "id": str(e.id), # Преобразуем UUID в строку
+ "name": e.name,
+ "type": e.type,
+ "has_embedding": embeddings[i] is not None,
+ "embedding_shape": str(embeddings[i].shape) if embeddings[i] is not None else None,
+ "text_length": len(e.text),
+ "in_search_text_length": len(e.in_search_text) if e.in_search_text else 0
+ }
+ for i, e in enumerate(entities[:5]) # Берем только первые 5 для примера
+ ]
+
+ return {
+ "stats": stats,
+ "examples": examples
+ }
\ No newline at end of file
diff --git a/routes/llm.py b/routes/llm.py
index a8cdc51f32e2ea62bf98ebfc9040f59033f3587b..0eb70596f4e5503d2ca0e0350e66de24b2752023 100644
--- a/routes/llm.py
+++ b/routes/llm.py
@@ -1,151 +1,129 @@
import logging
-from typing import Annotated, Optional, Tuple
import os
-from fastapi import APIRouter, BackgroundTasks, HTTPException, Response, UploadFile, Depends
-from components.llm.common import LlmParams, LlmPredictParams, Message
-from components.llm.deepinfra_api import DeepInfraApi
-from components.llm.llm_api import LlmApi
-from components.llm.common import ChatRequest
+from typing import Annotated, Optional
+from uuid import UUID
-from common.constants import PROMPT
-from components.llm.prompts import SYSTEM_PROMPT
-from components.llm.utils import append_llm_response_to_history, convert_to_openai_format
-from components.nmd.aggregate_answers import preprocessed_chunks
-from components.nmd.llm_chunk_search import LLMChunkSearch
from components.services.dataset import DatasetService
-from common.configuration import Configuration, Query, SummaryChunks
-from components.datasets.dispatcher import Dispatcher
-from common.exceptions import LLMResponseException
-from components.dbo.models.log import Log
+from components.services.entity import EntityService
+from fastapi import APIRouter, Depends, HTTPException
+
+import common.dependencies as DI
+from common.configuration import Configuration, Query
+from components.llm.common import ChatRequest, LlmParams, LlmPredictParams, Message
+from components.llm.deepinfra_api import DeepInfraApi
+from components.llm.utils import append_llm_response_to_history
from components.services.llm_config import LLMConfigService
from components.services.llm_prompt import LlmPromptService
-from schemas.dataset import (Dataset, DatasetExpanded, DatasetProcessing,
- SortQuery, SortQueryList)
-import common.dependencies as DI
-from sqlalchemy.orm import Session
router = APIRouter(prefix='/llm')
logger = logging.getLogger(__name__)
conf = DI.get_config()
-llm_params = LlmParams(**{
- "url": conf.llm_config.base_url,
- "model": conf.llm_config.model,
- "tokenizer": "unsloth/Llama-3.3-70B-Instruct",
- "type": "deepinfra",
- "default": True,
- "predict_params": LlmPredictParams(
- temperature=0.15, top_p=0.95, min_p=0.05, seed=42,
- repetition_penalty=1.2, presence_penalty=1.1, n_predict=2000
- ),
- "api_key": os.environ.get(conf.llm_config.api_key_env),
- "context_length": 128000
-})
-#TODO: унести в DI
+llm_params = LlmParams(
+ **{
+ "url": conf.llm_config.base_url,
+ "model": conf.llm_config.model,
+ "tokenizer": "unsloth/Llama-3.3-70B-Instruct",
+ "type": "deepinfra",
+ "default": True,
+ "predict_params": LlmPredictParams(
+ temperature=0.15,
+ top_p=0.95,
+ min_p=0.05,
+ seed=42,
+ repetition_penalty=1.2,
+ presence_penalty=1.1,
+ n_predict=2000,
+ ),
+ "api_key": os.environ.get(conf.llm_config.api_key_env),
+ "context_length": 128000,
+ }
+)
+# TODO: унести в DI
llm_api = DeepInfraApi(params=llm_params)
-@router.post("/chunks")
-def get_chunks(query: Query, dispatcher: Annotated[Dispatcher, Depends(DI.get_dispatcher)]) -> SummaryChunks:
- logger.info(f"Handling POST request to /chunks with query: {query.query}")
- try:
- result = dispatcher.search_answer(query)
- logger.info("Successfully retrieved chunks")
- return result
- except Exception as e:
- logger.error(f"Error retrieving chunks: {str(e)}")
- raise e
-
-
-def llm_answer(query: str, answer_chunks: SummaryChunks, config: Configuration
- ) -> Tuple[str, str, str, int]:
- """
- Метод для поиска правильного ответа с помощью LLM.
- Args:
- query: Запрос.
- answer_chunks: Ответы векторного поиска и elastic.
-
- Returns:
- Возвращает исходные chunks из поисков, и chunk который выбрала модель.
- """
- prompt = PROMPT
- llm_search = LLMChunkSearch(config.llm_config, PROMPT, logger)
- return llm_search.llm_chunk_search(query, answer_chunks, prompt)
-
-
-@router.post("/answer_llm")
-def get_llm_answer(query: Query, chunks: SummaryChunks, db: Annotated[Session, Depends(DI.get_db)], config: Annotated[Configuration, Depends(DI.get_config)]):
- logger.info(f"Handling POST request to /answer_llm with query: {query.query}")
- try:
- text_chunks, answer_llm, llm_prompt, _ = llm_answer(query.query, chunks, config)
-
- if not answer_llm:
- logger.error("LLM returned empty response")
- raise LLMResponseException()
-
- log_entry = Log(
- llmPrompt=llm_prompt,
- llmResponse=answer_llm,
- userRequest=query.query,
- query_type=chunks.query_type,
- userName=query.userName,
- )
- with db() as session:
- session.add(log_entry)
- session.commit()
- session.refresh(log_entry)
-
- logger.info(f"Successfully processed LLM request, log_id: {log_entry.id}")
- return {
- "answer_llm": answer_llm,
- "log_id": log_entry.id,
- }
-
- except Exception as e:
- logger.error(f"Error processing LLM request: {str(e)}")
- raise e
-
@router.post("/chat")
-async def chat(request: ChatRequest, config: Annotated[Configuration, Depends(DI.get_config)], llm_api: Annotated[DeepInfraApi, Depends(DI.get_llm_service)], prompt_service: Annotated[LlmPromptService, Depends(DI.get_llm_prompt_service)], llm_config_service: Annotated[LLMConfigService, Depends(DI.get_llm_config_service)], dispatcher: Annotated[Dispatcher, Depends(DI.get_dispatcher)]):
+async def chat(
+ request: ChatRequest,
+ config: Annotated[Configuration, Depends(DI.get_config)],
+ llm_api: Annotated[DeepInfraApi, Depends(DI.get_llm_service)],
+ prompt_service: Annotated[LlmPromptService, Depends(DI.get_llm_prompt_service)],
+ llm_config_service: Annotated[LLMConfigService, Depends(DI.get_llm_config_service)],
+ entity_service: Annotated[EntityService, Depends(DI.get_entity_service)],
+ dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)],
+):
try:
p = llm_config_service.get_default()
system_prompt = prompt_service.get_default()
-
+
predict_params = LlmPredictParams(
- temperature=p.temperature, top_p=p.top_p, min_p=p.min_p, seed=p.seed,
- frequency_penalty=p.frequency_penalty, presence_penalty=p.presence_penalty, n_predict=p.n_predict, stop=[]
+ temperature=p.temperature,
+ top_p=p.top_p,
+ min_p=p.min_p,
+ seed=p.seed,
+ frequency_penalty=p.frequency_penalty,
+ presence_penalty=p.presence_penalty,
+ n_predict=p.n_predict,
+ stop=[],
)
-
- #TODO: Вынести
+
+ # TODO: Вынести
def get_last_user_message(chat_request: ChatRequest) -> Optional[Message]:
return next(
(
- msg for msg in reversed(chat_request.history)
- if msg.role == "user" and (msg.searchResults is None or not msg.searchResults)
+ msg
+ for msg in reversed(chat_request.history)
+ if msg.role == "user"
+ and (msg.searchResults is None or not msg.searchResults)
),
- None
+ None,
)
-
- def insert_search_results_to_message(chat_request: ChatRequest, new_content: str) -> bool:
+
+ def insert_search_results_to_message(
+ chat_request: ChatRequest, new_content: str
+ ) -> bool:
for msg in reversed(chat_request.history):
- if msg.role == "user" and (msg.searchResults is None or not msg.searchResults):
+ if msg.role == "user" and (
+ msg.searchResults is None or not msg.searchResults
+ ):
msg.content = new_content
return True
return False
-
+
last_query = get_last_user_message(request)
search_result = None
+ logger.info(f"last_query: {last_query}")
+
if last_query:
- search_result = dispatcher.search_answer(Query(query=last_query.content, query_abbreviation=last_query.content))
- text_chunks = preprocessed_chunks(search_result, None, logger)
+ dataset = dataset_service.get_current_dataset()
+ if dataset is None:
+ raise HTTPException(status_code=400, detail="Dataset not found")
+ logger.info(f"last_query: {last_query.content}")
+ _, scores, chunk_ids = entity_service.search_similar(last_query.content, dataset.id)
+
+ chunks = entity_service.chunk_repository.get_chunks_by_ids(chunk_ids)
+
+ logger.info(f"chunk_ids: {chunk_ids[:3]}...{chunk_ids[-3:]}")
+ logger.info(f"scores: {scores[:3]}...{scores[-3:]}")
+
+ text_chunks = entity_service.build_text(chunks, scores)
+ logger.info(f"text_chunks: {text_chunks[:3]}...{text_chunks[-3:]}")
+
new_message = f'{last_query.content} /n/n{text_chunks}/n'
insert_search_results_to_message(request, new_message)
- response = await llm_api.predict_chat_stream(request, system_prompt.text, predict_params)
+ logger.info(f"request: {request}")
+
+ response = await llm_api.predict_chat_stream(
+ request, system_prompt.text, predict_params
+ )
result = append_llm_response_to_history(request, response)
return result
except Exception as e:
- logger.error(f"Error processing LLM request: {str(e)}", stack_info=True, stacklevel=10)
- return {"error": str(e)}
\ No newline at end of file
+ logger.error(
+ f"Error processing LLM request: {str(e)}", stack_info=True, stacklevel=10
+ )
+ return {"error": str(e)}
diff --git a/schemas/entity.py b/schemas/entity.py
new file mode 100644
index 0000000000000000000000000000000000000000..3775bd163099dbc3c6a21502c1c383c521b1539f
--- /dev/null
+++ b/schemas/entity.py
@@ -0,0 +1,58 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class EntitySearchRequest(BaseModel):
+ """Схема запроса для поиска сущностей."""
+ query: str
+ dataset_id: int
+
+
+class EntitySearchResponse(BaseModel):
+ """Схема ответа с результатами поиска сущностей."""
+ scores: List[float]
+ entity_ids: List[str]
+
+
+class EntitySearchWithTextRequest(BaseModel):
+ """Схема запроса для поиска сущностей с текстами."""
+ query: str
+ dataset_id: int
+
+
+class ChunkInfo(BaseModel):
+ """Информация о чанке."""
+ id: str
+ text: str
+ score: float
+
+
+class EntitySearchWithTextResponse(BaseModel):
+ """Схема ответа с результатами поиска сущностей и их текстами."""
+ chunks: List[ChunkInfo]
+
+
+class EntityTextRequest(BaseModel):
+ """Схема запроса для сборки текста из сущностей."""
+ entities: List[str]
+ chunk_scores: Optional[dict[str, float]] = None
+ include_tables: bool = True
+ max_documents: Optional[int] = None
+
+
+class EntityTextResponse(BaseModel):
+ """Схема ответа со сборкой текста из сущностей."""
+ text: str
+
+
+class EntityNeighborsRequest(BaseModel):
+ """Схема запроса для получения соседних чанков."""
+ entities: List[str]
+ max_distance: int = 1
+
+
+class EntityNeighborsResponse(BaseModel):
+ """Схема ответа с соседними чанками."""
+ entities: List[str]
+
\ No newline at end of file
diff --git a/scripts/analyze_entities.py b/scripts/analyze_entities.py
new file mode 100644
index 0000000000000000000000000000000000000000..11e701fd37bf9dbc5108b89932cc3e395c30b982
--- /dev/null
+++ b/scripts/analyze_entities.py
@@ -0,0 +1,150 @@
+import argparse
+import logging
+from typing import Optional
+
+import numpy as np
+from sqlalchemy.orm import Session
+
+import common.dependencies as DI
+from common.configuration import Configuration
+from components.dbo.models.entity import EntityModel
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def analyze_embeddings(embeddings: list[Optional[np.ndarray]]) -> dict:
+ """
+ Анализ эмбеддингов.
+
+ Args:
+ embeddings: Список эмбеддингов
+
+ Returns:
+ dict: Статистика по эмбеддингам
+ """
+ valid_embeddings = [e for e in embeddings if e is not None]
+ if not valid_embeddings:
+ return {
+ "total": len(embeddings),
+ "valid": 0,
+ "shapes": {},
+ "mean_norm": None,
+ "std_norm": None
+ }
+
+ shapes = {}
+ norms = []
+ for e in valid_embeddings:
+ shape_str = str(e.shape)
+ shapes[shape_str] = shapes.get(shape_str, 0) + 1
+ norms.append(np.linalg.norm(e))
+
+ return {
+ "total": len(embeddings),
+ "valid": len(valid_embeddings),
+ "shapes": shapes,
+ "mean_norm": float(np.mean(norms)),
+ "std_norm": float(np.std(norms))
+ }
+
+
+def analyze_entities(
+ dataset_id: int,
+ db: Session,
+ config: Configuration,
+) -> None:
+ """
+ Анализ сущностей в датасете.
+
+ Args:
+ dataset_id: ID датасета
+ db: Сессия базы данных
+ config: Конфигурация приложения
+ """
+ # Получаем все сущности
+ entities = (
+ db.query(EntityModel)
+ .filter(EntityModel.dataset_id == dataset_id)
+ .all()
+ )
+
+ if not entities:
+ logger.error(f"No entities found for dataset {dataset_id}")
+ return
+
+ # Базовая статистика
+ logger.info(f"Total entities: {len(entities)}")
+ logger.info(f"Entity types: {set(e.entity_type for e in entities)}")
+
+ # Статистика по типам
+ type_stats = {}
+ for e in entities:
+ if e.entity_type not in type_stats:
+ type_stats[e.entity_type] = 0
+ type_stats[e.entity_type] += 1
+
+ logger.info("Entities per type:")
+ for t, count in type_stats.items():
+ logger.info(f" {t}: {count}")
+
+ # Анализ эмбеддингов
+ embeddings = [e.embedding for e in entities]
+ embedding_stats = analyze_embeddings(embeddings)
+
+ logger.info("\nEmbedding statistics:")
+ logger.info(f" Total embeddings: {embedding_stats['total']}")
+ logger.info(f" Valid embeddings: {embedding_stats['valid']}")
+ logger.info(" Shapes:")
+ for shape, count in embedding_stats['shapes'].items():
+ logger.info(f" {shape}: {count}")
+ if embedding_stats['mean_norm'] is not None:
+ logger.info(f" Mean norm: {embedding_stats['mean_norm']:.4f}")
+ logger.info(f" Std norm: {embedding_stats['std_norm']:.4f}")
+
+ # Анализ текстов
+ text_lengths = [len(e.text) for e in entities]
+ search_text_lengths = [len(e.in_search_text) if e.in_search_text else 0 for e in entities]
+
+ logger.info("\nText statistics:")
+ logger.info(f" Mean text length: {np.mean(text_lengths):.2f}")
+ logger.info(f" Std text length: {np.std(text_lengths):.2f}")
+ logger.info(f" Mean search text length: {np.mean(search_text_lengths):.2f}")
+ logger.info(f" Std search text length: {np.std(search_text_lengths):.2f}")
+
+ # Примеры сущностей
+ logger.info("\nExample entities:")
+ for e in entities[:5]:
+ logger.info(f" ID: {e.uuid}")
+ logger.info(f" Name: {e.name}")
+ logger.info(f" Type: {e.entity_type}")
+ logger.info(f" Embedding: {e.embedding}")
+ if e.embedding is not None:
+ logger.info(f" Embedding shape: {e.embedding.shape}")
+ logger.info(" ---")
+
+
+def main() -> None:
+ """Точка входа скрипта."""
+ parser = argparse.ArgumentParser(description="Analyze entities in dataset")
+ parser.add_argument("dataset_id", type=int, help="Dataset ID")
+ parser.add_argument(
+ "--config",
+ type=str,
+ default="config_dev.yaml",
+ help="Path to config file",
+ )
+ args = parser.parse_args()
+
+ config = Configuration(args.config)
+ db = DI.get_db()
+
+ with db() as session:
+ try:
+ analyze_entities(args.dataset_id, session, config)
+ finally:
+ session.close()
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file