from unsloth import FastVisionModel
from consts import BASE_MODEL


def setup_model(model: str) -> tuple:
    model, tokenizer = FastVisionModel.from_pretrained(
        BASE_MODEL,
        load_in_4bit = True, 
        use_gradient_checkpointing = "True",
    )
    
    model = FastVisionModel.get_peft_model(
        model,
        finetune_vision_layers     = False, 
        finetune_language_layers   = True, 
        finetune_attention_modules = True, 
        finetune_mlp_modules       = True,

        r = 16, 
        lora_alpha = 16,  
        lora_dropout = 0,
        bias = "none",
        random_state = 3407,
        use_rslora = False, 
        loftq_config = None, 
        use_gradient_checkpointing = "True" 
    )

    return model, tokenizer