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)
    # question_crawling="나의 이런 점은 일할 때 도움이 돼!?"
    question_crawling="즐겨보는 유튜버?"
    # answer_crawling="나누고 싶어하는 마음? 주변 사람들은 그만 퍼주라고 하기도 하지만, 내가 클라이언트로부터 돈을 벌고자 하는 것이 아니라 조금이라도 더 챙겨주고자 하는 마음을 가졌을 때 결국 나의 브랜드가 훨씬 더 커질 수 있다는 믿음이 있다."
    answer_crawling="선바, 드립치는 게 제 취향입니다."

    input_text = f"""
    <instruction>
    Using the text: {question_crawling} {answer_crawling}, create a new multiple-choice question with 4 answer options.
    """

    print(model.generate(input_text, max_length=512))

    return model


# def template_setting(df:pd.DataFrame, is_test:bool) -> np.ndarray:
#     template_input="""
#     <instruction>
#     Using the text: {question_crawling} {answer_crawling}, create a new multiple-choice question with 4 answer options.
#     """
    
#     template_output="""
#     <Response>
#     {question_generated}
#     {multiple_choice_generated}
#     {answer_generated}
#     """
#     template=template_input+'\n'+template_output
    
#     inputs = np.array(df.apply(lambda row: template.format(
#                         question_crawling=row['question_crawling'],
#                         answer_crawling=row['answer_crawling'],
#                         question_generated=row['question_generated'] if not is_test else "",
#                         multiple_choice_generated=row['multiple_choice_generated'] if not is_test else "",
#                         answer_generated=row['answer_generated'] if not is_test else "").strip(), axis=1))
    
#     outputs = np.array(df.apply(lambda row: template_output.format(
#                         question_generated=row['question_generated'],
#                         multiple_choice_generated=row['multiple_choice_generated'],
#                         answer_generated=row['answer_generated']).strip(), axis=1))
#     combined_array = np.column_stack((inputs, outputs))
#     return combined_array