configuration_magiv2.py

from transformers import PretrainedConfig, VisionEncoderDecoderConfig
from typing import Any, Optional


class Magiv2Config(PretrainedConfig):
    """
    Klasa konfiguracyjna dla modelu Magiv2.

    Magiv2Config dziedziczy po PretrainedConfig z biblioteki transformers i definiuje
    kompletną konfigurację dla modelu wizyjnego składającego się z trzech głównych komponentów:
    - Model detekcji obiektów (detection)
    - Model OCR (rozpoznawanie tekstu)
    - Model embedowania wyciętych fragmentów obrazu (crop embeddings)

    Attributes:
        model_type: Identyfikator typu modelu dla biblioteki transformers
        disable_ocr: Flaga wyłączająca moduł OCR
        disable_crop_embeddings: Flaga wyłączająca moduł embedowania wyciętych fragmentów
        disable_detections: Flaga wyłączająca moduł detekcji obiektów
        detection_model_config: Konfiguracja modelu detekcji (po deserializacji)
        ocr_model_config: Konfiguracja modelu OCR (po deserializacji)
        crop_embedding_model_config: Konfiguracja modelu embedowania (po deserializacji)
        detection_image_preprocessing_config: Parametry przetwarzania obrazu dla detekcji
        ocr_pretrained_processor_path: Ścieżka do wytrenowanego procesora OCR
        crop_embedding_image_preprocessing_config: Parametry przetwarzania obrazu dla embedowania
    """

    # Identyfikator typu modelu używany przez bibliotekę transformers
    model_type: str = "magiv2"

    def __init__(
        self,
        disable_ocr: bool = False,
        disable_crop_embeddings: bool = False,
        disable_detections: bool = False,
        detection_model_config: Optional[dict[str, Any]] = None,
        ocr_model_config: Optional[dict[str, Any]] = None,
        crop_embedding_model_config: Optional[dict[str, Any]] = None,
        detection_image_preprocessing_config: Optional[dict[str, Any]] = None,
        ocr_pretrained_processor_path: Optional[str] = None,
        crop_embedding_image_preprocessing_config: Optional[dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """
        Inicjalizuje konfigurację modelu Magiv2.

        Konstruktor przyjmuje parametry kontrolujące które moduły modelu są aktywne,
        oraz konfiguracje dla poszczególnych komponentów. Konfiguracje przekazane jako
        słowniki są deserializowane do odpowiednich obiektów Config z transformers.

        Args:
            disable_ocr: Czy wyłączyć moduł rozpoznawania tekstu (OCR).
                        Domyślnie False - OCR jest aktywne.
            disable_crop_embeddings: Czy wyłączyć moduł tworzenia embeddingów dla wyciętych
                                    fragmentów obrazu. Domyślnie False - embedowanie aktywne.
            disable_detections: Czy wyłączyć moduł detekcji obiektów na obrazie.
                               Domyślnie False - detekcja aktywna.
            detection_model_config: Słownik z konfiguracją modelu detekcji obiektów.
                                   Jeśli podany, zostanie zdeserializowany do PretrainedConfig.
            ocr_model_config: Słownik z konfiguracją modelu OCR (encoder-decoder).
                             Jeśli podany, zostanie zdeserializowany do VisionEncoderDecoderConfig.
            crop_embedding_model_config: Słownik z konfiguracją modelu embedowania wyciętych
                                        fragmentów. Jeśli podany, zostanie zdeserializowany
                                        do PretrainedConfig.
            detection_image_preprocessing_config: Słownik z parametrami preprocessingu obrazu
                                                 dla modułu detekcji (np. rozmiar, normalizacja).
            ocr_pretrained_processor_path: Ścieżka do katalogu lub Hub ID z wytrenowanym
                                          procesorem obrazu dla modułu OCR.
            crop_embedding_image_preprocessing_config: Słownik z parametrami preprocessingu
                                                      obrazu dla modułu embedowania.
            **kwargs: Dodatkowe argumenty przekazywane do klasy bazowej PretrainedConfig.

        Returns:
            None

        Note:
            - Konfiguracje modeli są deserializowane z dict do obiektów Config tylko wtedy,
              gdy zostały przekazane (nie są None)
            - Flagi disable_* pozwalają na selektywne wyłączanie poszczególnych modułów
            - Wszystkie dodatkowe kwargs są przekazywane do klasy bazowej PretrainedConfig
        """
        # Przechowywanie flag wyłączających poszczególne moduły
        self.disable_ocr: bool = disable_ocr
        self.disable_crop_embeddings: bool = disable_crop_embeddings
        self.disable_detections: bool = disable_detections

        # Przechowywanie dodatkowych argumentów przekazanych do konstruktora
        self.kwargs: dict[str, Any] = kwargs

        # Inicjalizacja atrybutów konfiguracji modeli jako None
        # (mogą zostać zdeserializowane poniżej jeśli parametry nie są None)
        self.detection_model_config: Optional[PretrainedConfig] = None
        self.ocr_model_config: Optional[VisionEncoderDecoderConfig] = None
        self.crop_embedding_model_config: Optional[PretrainedConfig] = None

        # Deserializacja konfiguracji modelu detekcji ze słownika do obiektu PretrainedConfig
        if detection_model_config is not None:
            self.detection_model_config = PretrainedConfig.from_dict(
                detection_model_config
            )

        # Deserializacja konfiguracji modelu OCR ze słownika do obiektu VisionEncoderDecoderConfig
        # OCR wykorzystuje architekturę encoder-decoder (vision encoder + text decoder)
        if ocr_model_config is not None:
            self.ocr_model_config = VisionEncoderDecoderConfig.from_dict(
                ocr_model_config
            )

        # Deserializacja konfiguracji modelu embedowania ze słownika do obiektu PretrainedConfig
        if crop_embedding_model_config is not None:
            self.crop_embedding_model_config = PretrainedConfig.from_dict(
                crop_embedding_model_config
            )

        # Przechowywanie konfiguracji preprocessingu obrazu dla modułu detekcji
        # (np. docelowy rozmiar obrazu, parametry normalizacji, augmentacje)
        self.detection_image_preprocessing_config: Optional[dict[str, Any]] = (
            detection_image_preprocessing_config
        )

        # Ścieżka do wytrenowanego procesora OCR (może być lokalna lub z Hugging Face Hub)
        self.ocr_pretrained_processor_path: Optional[str] = ocr_pretrained_processor_path

        # Przechowywanie konfiguracji preprocessingu obrazu dla modułu embedowania
        # (np. docelowy rozmiar wycięć, parametry normalizacji)
        self.crop_embedding_image_preprocessing_config: Optional[dict[str, Any]] = (
            crop_embedding_image_preprocessing_config
        )

        # Wywołanie konstruktora klasy bazowej PretrainedConfig z dodatkowymi kwargs
        super().__init__(**kwargs)