import keras_nlp

MODEL_NAME = "gemma2_instruct_2b_en"
LORA_WEIGHT_PATH = "ice_breaking_challenge/models/gemma2_it_2b_icebreaking.lora.h5"

def load_model_with_lora(model_name:str = MODEL_NAME, lora_weight_path: str = LORA_WEIGHT_PATH):
    """
    Keras 기반 모델 로드 및 LoRA 가중치 적용

    Args:
        model_name (str): 로드할 모델의 이름
        lora_weight_path (str): 적용할 LoRA 가중치 파일의 경로

    Returns:
        keras_nlp.models.GemmaCausalLM: 로드된 모델
    """
    model = keras_nlp.models.GemmaCausalLM.from_preset(model_name)

    model.backbone.load_lora_weights(lora_weight_path)

    return model