magiv2 / configuration_magiv2.py
ragavsachdeva's picture
Upload model
f7499c0 verified
raw
history blame contribute delete
No virus
1.7 kB
from transformers import PretrainedConfig, VisionEncoderDecoderConfig
from typing import List
class Magiv2Config(PretrainedConfig):
model_type = "magiv2"
def __init__(
self,
disable_ocr: bool = False,
disable_crop_embeddings: bool = False,
disable_detections: bool = False,
detection_model_config: dict = None,
ocr_model_config: dict = None,
crop_embedding_model_config: dict = None,
detection_image_preprocessing_config: dict = None,
ocr_pretrained_processor_path: str = None,
crop_embedding_image_preprocessing_config: dict = None,
**kwargs,
):
self.disable_ocr = disable_ocr
self.disable_crop_embeddings = disable_crop_embeddings
self.disable_detections = disable_detections
self.kwargs = kwargs
self.detection_model_config = None
self.ocr_model_config = None
self.crop_embedding_model_config = None
if detection_model_config is not None:
self.detection_model_config = PretrainedConfig.from_dict(detection_model_config)
if ocr_model_config is not None:
self.ocr_model_config = VisionEncoderDecoderConfig.from_dict(ocr_model_config)
if crop_embedding_model_config is not None:
self.crop_embedding_model_config = PretrainedConfig.from_dict(crop_embedding_model_config)
self.detection_image_preprocessing_config = detection_image_preprocessing_config
self.ocr_pretrained_processor_path = ocr_pretrained_processor_path
self.crop_embedding_image_preprocessing_config = crop_embedding_image_preprocessing_config
super().__init__(**kwargs)