from dataclasses import dataclass, field
from typing import List, Optional
from transformers import TrainingArguments
from src.models.ret2.configuration_ret2 import RetLayerStrategy, LOGIT_SCALE_CLIP_PRETRAINED
from src import custom_utils

logger = custom_utils.get_logger(__name__)

@dataclass
class CustomTrainingArguments(TrainingArguments):
    dataset_n_samples: Optional[int] = None
    dataset_eval_n_samples: int = 10000
    dataset_train: List[str] = field(default_factory=lambda: [])
    dataset_eval: List[str] = field(default_factory=lambda: [])
    dataset_ratio: List[float] = field(default_factory=lambda: [])
    dataset_sample_mbeir_instructions: bool = False
    
    # https://huggingface.co/datasets/TIGER-Lab/M-BEIR/blob/main/instructions/query_instructions.tsv
    dataset_mbeir_instructions_path: str = 'M-BEIR/instructions/query_instructions.tsv'

    mbeir_datatask_balance_sampler: bool = False
    eval_similarity_matrix_batch_size: Optional[int] = 128
    backbone_learning_rate_multiplier: Optional[float] = None
    do_only_eval : Optional[bool] = False
    resume_from_last_checkpoint: Optional[bool] = False


@dataclass
class ModelArguments:
    text_name_or_path: str = 'openai/clip-vit-large-patch14'
    text_frozen: bool = True
    text_layer_strategy: str = RetLayerStrategy.VIT_B
    text_global_feats_proj_size: int = 0
    vision_name_or_path: str = 'openai/clip-vit-large-patch14'
    vision_frozen: bool = True
    vision_layer_strategy: str = RetLayerStrategy.VIT_L
    vision_global_feats_proj_size: int = 0
    tokenizer_max_length: int = None
    w_final_size: int = 768
    logit_scale_init_value: float = LOGIT_SCALE_CLIP_PRETRAINED
    logit_scale_clamp: float = 100.0
    hidden_size: int = 1024
    dropout_p: float = 0.05
    attention_dropout: float = 0.05
    input_gate_bias: float = 0.0
    forget_gate_bias: float = 0.0         
