|
|
import sys
|
|
|
from pathlib import Path
|
|
|
parent_root = Path().resolve().parent.parent
|
|
|
sys.path.append(str(parent_root))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from transformers import PretrainedConfig, DonutSwinConfig, GemmaConfig, CONFIG_MAPPING, SiglipVisionConfig
|
|
|
from typing import Tuple, Literal
|
|
|
|
|
|
|
|
|
|
|
|
class PamConfig(PretrainedConfig):
|
|
|
model_type = "pam"
|
|
|
def __init__(
|
|
|
self,
|
|
|
sequence_mapping_layer_type: Literal["linear_projection","bilinear_interpolation"] = "bilinear_interpolation",
|
|
|
student_fmap_dim: Tuple[int,int]=(80,60),
|
|
|
student_embedding_dim: int = 1024,
|
|
|
teacher_fmap_dim: Tuple[int,int] = (64,64),
|
|
|
teacher_embedding_dim: int = 1152,
|
|
|
**kwargs,
|
|
|
):
|
|
|
self.sequence_mapping_layer_type = sequence_mapping_layer_type
|
|
|
self.student_fmap_dim = student_fmap_dim
|
|
|
self.student_embedding_dim = student_embedding_dim
|
|
|
self.teacher_fmap_dim = teacher_fmap_dim
|
|
|
self.teacher_embedding_dim = teacher_embedding_dim
|
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
|
|
|
class SwinPamVisionEncoderConfig(PretrainedConfig):
|
|
|
model_type = "swinpam"
|
|
|
sub_configs = {"encoder_config": DonutSwinConfig, "pam_config": PamConfig}
|
|
|
def __init__(
|
|
|
self,
|
|
|
encoder_config: DonutSwinConfig = None,
|
|
|
pam_config: PamConfig = None,
|
|
|
**kwargs
|
|
|
):
|
|
|
self.encoder_config = encoder_config
|
|
|
self.pam_config = pam_config
|
|
|
|
|
|
if isinstance(self.encoder_config, dict):
|
|
|
encoder_config["model_type"] = (
|
|
|
encoder_config["model_type"] if "model_type" in encoder_config else "donut-swin"
|
|
|
)
|
|
|
if encoder_config["model_type"] == "donut-swin":
|
|
|
self.encoder_config = DonutSwinConfig(**encoder_config)
|
|
|
else:
|
|
|
print(f"Encoder type: {encoder_config['model_type']}")
|
|
|
self.encoder_config = CONFIG_MAPPING[encoder_config["model_type"]](**encoder_config)
|
|
|
|
|
|
'''
|
|
|
elif encoder_config is None:
|
|
|
print("coucou2")
|
|
|
self.encoder_config = DonutSwinConfig()
|
|
|
'''
|
|
|
|
|
|
if isinstance(self.pam_config, dict):
|
|
|
'''
|
|
|
pam_config["model_type"] = (
|
|
|
pam_config["model_type"] if "model_type" in pam_config else "pam"
|
|
|
)
|
|
|
'''
|
|
|
if pam_config["model_type"] == "pam":
|
|
|
self.pam_config = PamConfig(**pam_config)
|
|
|
else:
|
|
|
raise ValueError(f"pam_config['model_type'] should be 'pam', got {pam_config['model_type']}")
|
|
|
'''
|
|
|
elif pam_config is None:
|
|
|
self.pam_config = PamConfig()
|
|
|
'''
|
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
|
|
|
class SiglipPAMVisionEncoderConfig(PretrainedConfig):
|
|
|
model_type = "siglippam"
|
|
|
sub_configs = {"encoder_config": SiglipVisionConfig, "pam_config": PamConfig}
|
|
|
def __init__(
|
|
|
self,
|
|
|
encoder_config: SiglipVisionConfig = None,
|
|
|
pam_config: PamConfig = None,
|
|
|
**kwargs
|
|
|
):
|
|
|
self.encoder_config = encoder_config
|
|
|
self.pam_config = pam_config
|
|
|
|
|
|
if isinstance(self.encoder_config, dict):
|
|
|
encoder_config["model_type"] = (
|
|
|
encoder_config["model_type"] if "model_type" in encoder_config else "siglip_vision_model"
|
|
|
)
|
|
|
if encoder_config["model_type"] == "siglip_vision_model":
|
|
|
self.encoder_config = SiglipVisionConfig(**encoder_config)
|
|
|
else:
|
|
|
raise ValueError(f"Need siglip_model_type, got {encoder_config['model_type']}")
|
|
|
|
|
|
if isinstance(self.pam_config, dict):
|
|
|
if pam_config["model_type"] == "pam":
|
|
|
self.pam_config = PamConfig(**pam_config)
|
|
|
else:
|
|
|
raise ValueError(f"pam_config['model_type'] should be 'pam', got {pam_config['model_type']}")
|
|
|
|
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
|
|
|
class DIVEdocConfig(PretrainedConfig):
|
|
|
keys_to_ignore_at_inference = ["past_key_values"]
|
|
|
sub_configs = {"vision_config": SwinPamVisionEncoderConfig, "text_config": GemmaConfig}
|
|
|
model_type = "DIVEdoc"
|
|
|
def __init__(
|
|
|
self,
|
|
|
vision_config=None,
|
|
|
text_config=None,
|
|
|
ignore_index=-100,
|
|
|
image_token_index=256000,
|
|
|
vocab_size=257152,
|
|
|
projection_dim=2048,
|
|
|
hidden_size=2048,
|
|
|
|
|
|
**kwargs,
|
|
|
):
|
|
|
self._ignore_index = ignore_index
|
|
|
self.image_token_index = image_token_index
|
|
|
self._vocab_size = vocab_size
|
|
|
self.projection_dim = projection_dim
|
|
|
self.hidden_size = hidden_size
|
|
|
self.vision_config = vision_config
|
|
|
self.is_encoder_decoder = False
|
|
|
|
|
|
|
|
|
|
|
|
if isinstance(self.vision_config, dict):
|
|
|
vision_config["model_type"] = (
|
|
|
vision_config["model_type"] if "model_type" in vision_config else "swinpam"
|
|
|
)
|
|
|
if vision_config["model_type"] == "swinpam":
|
|
|
self.vision_config = SwinPamVisionEncoderConfig(encoder_config=vision_config["encoder_config"],pam_config=vision_config["pam_config"])
|
|
|
elif vision_config["model_type"] == "siglippam":
|
|
|
self.vision_config = SiglipPAMVisionEncoderConfig(encoder_config=vision_config["encoder_config"],pam_config=vision_config["pam_config"])
|
|
|
else:
|
|
|
self.vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
|
|
|
elif vision_config is None:
|
|
|
self.vision_config = get_vision_config("swinpam")
|
|
|
|
|
|
self.text_config = text_config
|
|
|
if isinstance(self.text_config, dict):
|
|
|
text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "gemma"
|
|
|
self.text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
|
|
|
elif text_config is None:
|
|
|
self.text_config = CONFIG_MAPPING["gemma"](
|
|
|
hidden_size=2048,
|
|
|
num_hidden_layers=18,
|
|
|
intermediate_size=16384,
|
|
|
num_attention_heads=8,
|
|
|
num_key_value_heads=1,
|
|
|
is_encoder_decoder=False,
|
|
|
vocab_size=vocab_size,
|
|
|
)
|
|
|
|
|
|
self.text_config.num_image_tokens = self.vision_config.pam_config.teacher_fmap_dim[0] *\
|
|
|
self.vision_config.pam_config.teacher_fmap_dim[1]
|
|
|
self.vision_config.projection_dim = projection_dim
|
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
def to_dict(self):
|
|
|
output = super().to_dict()
|
|
|
output.pop("_ignore_index", None)
|
|
|
return output
|
|
|
|
|
|
def get_siglip_vision_config(image_size=[896,896],num_image_token = 4096,hidden_size = 768):
|
|
|
encoder_config = SiglipVisionConfig(
|
|
|
hidden_size = hidden_size,
|
|
|
image_size = image_size,
|
|
|
intermediate_size = 2860,
|
|
|
model_type = "siglip_vision_model",
|
|
|
num_attention_heads = 8,
|
|
|
num_hidden_layers = 12,
|
|
|
num_image_tokens = num_image_token,
|
|
|
patch_size = 14,
|
|
|
projection_dim = 2048,
|
|
|
projector_hidden_act = "gelu_fast",
|
|
|
torch_dtype = "float32",
|
|
|
vision_use_head = False
|
|
|
)
|
|
|
return encoder_config
|
|
|
|
|
|
def get_swin_vision_config(image_size=[2560,1920],hidden_size = 1024):
|
|
|
encoder_config = DonutSwinConfig(
|
|
|
attention_probs_dropout_prob= 0.0,
|
|
|
depths =[
|
|
|
2,
|
|
|
2,
|
|
|
14,
|
|
|
2
|
|
|
],
|
|
|
drop_path_rate= 0.1,
|
|
|
embed_dim =128,
|
|
|
hidden_act ="gelu",
|
|
|
hidden_dropout_prob = 0.0,
|
|
|
hidden_size = hidden_size,
|
|
|
image_size = image_size,
|
|
|
initializer_range = 0.02,
|
|
|
layer_norm_eps = 1e-05,
|
|
|
mlp_ratio = 4.0,
|
|
|
model_type = "donut-swin",
|
|
|
num_channels = 3,
|
|
|
num_heads =[
|
|
|
4,
|
|
|
8,
|
|
|
16,
|
|
|
32
|
|
|
],
|
|
|
num_layers =4,
|
|
|
patch_size = 4,
|
|
|
path_norm = True,
|
|
|
qkv_bias = True,
|
|
|
use_absolute_embeddings = False,
|
|
|
window_size = 10
|
|
|
)
|
|
|
return encoder_config
|
|
|
|
|
|
def get_vision_config( visual_encoder_type:Literal["swinpam","siglip80m"],
|
|
|
image_size=[2560,1920],
|
|
|
sequence_mapping_layer_type= "bilinear",
|
|
|
student_fmap_dim=(80,60),
|
|
|
student_embedding_dim= 1024,
|
|
|
teacher_fmap_dim= (64,64),
|
|
|
teacher_embedding_dim= 1152):
|
|
|
pam_config = PamConfig(
|
|
|
sequence_mapping_layer_type = sequence_mapping_layer_type,
|
|
|
student_fmap_dim = student_fmap_dim,
|
|
|
student_embedding_dim = student_embedding_dim,
|
|
|
teacher_fmap_dim = teacher_fmap_dim,
|
|
|
teacher_embedding_dim = teacher_embedding_dim)
|
|
|
|
|
|
if visual_encoder_type == "swinpam":
|
|
|
encoder_config = get_swin_vision_config(image_size=image_size,hidden_size = student_embedding_dim)
|
|
|
ve_config = SwinPamVisionEncoderConfig(encoder_config=encoder_config,pam_config=pam_config)
|
|
|
return ve_config
|
|
|
|
|
|
elif visual_encoder_type =="siglip80m":
|
|
|
encoder_config = get_siglip_vision_config(image_size=image_size,num_image_token = (image_size//14)**2, hidden_size = student_embedding_dim)
|
|
|
ve_config = SiglipPAMVisionEncoderConfig(encoder_config=encoder_config,pam_config=pam_config)
|
|
|
return ve_config
|
|
|
else:
|
|
|
raise ValueError(f"Unknown visual encoder type, need 'swinpam' or 'siglip80m, got {visual_encoder_type}.") |
