#!/usr/bin/env python
# coding=utf-8
# Copyright 2023 Statistics and Machine Learning Research Group at HKUST. All rights reserved.
"""A one-line summary of the module or program, terminated by a period.

Leave one blank line.  The rest of this docstring should contain an
overall description of the module or program.  Optionally, it may also
contain a brief description of exported classes and functions and/or usage
examples.

Typical usage example:

  foo = ClassFoo()
  bar = foo.FunctionBar()
"""

import sys
import os

import numpy as np
import torch.serialization

# PyTorch 2.6 defaults torch.load(weights_only=True); allow Trainer RNG state pickles that
# rely on numpy._core.multiarray._reconstruct from checkpoints we produced ourselves.
torch.serialization.add_safe_globals([
    np.dtype,
    np.ufunc,
    (np.core.multiarray._reconstruct, "numpy._core.multiarray._reconstruct"),
    np.core.multiarray._reconstruct,
])
# sys.path.remove(os.path.abspath(os.path.dirname(sys.argv[0])))
from transformers import HfArgumentParser

from lmflow.args import (
    ModelArguments,
    DatasetArguments,
    AutoArguments,
)

from lmflow.datasets.dataset import Dataset
from lmflow.models.auto_model import AutoModel
from lmflow.pipeline.auto_pipeline import AutoPipeline


def main():
	# Parses arguments
    pipeline_name = "finetuner"
    PipelineArguments = AutoArguments.get_pipeline_args_class(pipeline_name)

    parser = HfArgumentParser((ModelArguments, DatasetArguments, PipelineArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, pipeline_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, pipeline_args = parser.parse_args_into_dataclasses()

    # Initialization
    finetuner = AutoPipeline.get_pipeline(
        pipeline_name=pipeline_name,
        model_args=model_args,
        data_args=data_args,
        pipeline_args=pipeline_args,
    )
    dataset = Dataset(data_args)
    model = AutoModel.get_model(model_args)

    if model_args.mdm:
        model.backend_model.bd_size = model_args.bd_size
        model.backend_model.config.bd_size = model_args.bd_size

        # if model_args.mdm_loss:
        #     model.backend_model.loss_type = "ForMaskedLM"
            
        # model.backend_model.config.mdm = model_args.mdm
        # model.backend_model.model.mdm = model_args.mdm
        model.backend_model.model.mask_id = model.tokenizer.encode("|<MASK>|")[0]

        model.backend_model.use_block_causal_mask = model_args.use_block_causal_mask
        model.backend_model.config.use_block_causal_mask = model_args.use_block_causal_mask

        model.backend_model.always_mask_im_end = model_args.always_mask_im_end
        model.backend_model.config.always_mask_im_end = model_args.always_mask_im_end

        model.backend_model.complementary_mask = model_args.complementary_mask
        model.backend_model.config.complementary_mask = model_args.complementary_mask

        model.backend_model.block_causal_no_dynamic = model_args.block_causal_no_dynamic
        model.backend_model.config.block_causal_no_dynamic = model_args.block_causal_no_dynamic

        data_args.bd_size = model_args.bd_size

    # model.backend_model.config.eos_loss_weight = model_args.eos_loss_weight
    # model.backend_model.eos_loss_weight = model_args.eos_loss_weight

    # if not model_args.mdm:
    # for n, param in model.backend_model.named_parameters():
    #     if 'embed_tokens' in n:
    #         param.requires_grad = True
    #     else:
    #         param.requires_grad = False

    # Finetuning
    tuned_model = finetuner.tune(model=model, dataset=dataset)


if __name__ == '__main__':
    main()
