#  Copyright (c) 2023 Intel Corporation
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
from pathlib import Path
import argparse
from typing import List, Optional
from transformers import AutoConfig
import subprocess
from huggingface_hub import snapshot_download

model_maps = {"gpt_neox": "gptneox", "gpt_bigcode": "starcoder"}
build_path = Path(Path(__file__).parent.absolute(), "../build/")


def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


def main(args_in: Optional[List[str]] = None) -> None:
    parser = argparse.ArgumentParser(description="run quantization and inference")
    parser.add_argument("model", type=Path, help="directory containing model file or model id")
    parser.add_argument("--build_dir", type=Path, help="path to build directory", default=build_path)

    # quantization related arguments.
    parser.add_argument(
        "--weight_dtype",
        choices=["int4", "int8", "fp8", "fp8_e5m2", "fp8_e4m3", "fp4", "fp4_e2m1", "nf4"],
        help="Data type of quantized weight: int4/int8 (default int4)",
        default="int4",
    )
    parser.add_argument(
        "--model_type",
        type=str,
        help="Setting the model_type manually)",
        default=None,
    )
    parser.add_argument(
        "--alg",
        type=str,
        choices=["sym", "asym"],
        help="Quantization algorithm: sym/asym (default sym)",
        default="sym",
    )
    parser.add_argument(
        "--group_size",
        type=int,
        choices=[-1, 32, 128],
        help="Group size: Int (default: 32)",
        default=32,
    )
    parser.add_argument(
        "--scale_dtype",
        type=str,
        choices=["fp32", "bf16", "fp8"],
        help="Data type of scales: fp32/bf16 (default fp32)",
        default="fp32",
    )
    parser.add_argument(
        "--compute_dtype",
        type=str,
        choices=["fp32", "fp16", "bf16", "int8"],
        help="Data type of Gemm computation: int8/bf16/fp32 (default: int8)",
        default="int8",
    )
    parser.add_argument("--format",
                        type=str,
                        default="NE",
                        choices=["NE", "GGUF"],
                        help="Convert to the GGUF or NE format")
    parser.add_argument(
        "--use_ggml",
        action="store_true",
        help="enable ggml for quantization and inference",
    )
    # inference related arguments.
    parser.add_argument(
        "-p",
        "--prompt",
        type=str,
        help="Prompt to start generation with: String (default: empty)",
        default="Once upon a time, there existed a ",
    )
    parser.add_argument(
        "-f",
        "--file",
        type=str,
        help="Path to a text file containing the prompt (for large prompts)",
        default=None,
    )
    parser.add_argument(
        "-n",
        "--n_predict",
        type=int,
        help="Number of tokens to predict: Int (default: -1, -1 = infinity)",
        default=-1,
    )
    parser.add_argument(
        "-t",
        "--threads",
        type=int,
        help="Number of threads to use during computation: Int (default: 56)",
        default=56,
    )
    parser.add_argument(
        "-b",
        "--batch_size_truncate",
        type=int,
        help="Batch size for prompt processing: Int (default: 512)",
        default=512,
    )
    parser.add_argument(
        "-c",
        "--ctx_size",
        type=int,
        help=
        "Size of the prompt context: Int (default: 512, can not be larger than specific model's context window length)",
        default=512,
    )
    parser.add_argument(
        "-s",
        "--seed",
        type=int,
        help="NG seed: Int (default: -1, use random seed for < 0)",
        default=-1,
    )
    parser.add_argument(
        "--repeat_penalty",
        type=float,
        help="NG seed: Int (default: -1, use random seed for < 0)",
        default=1.1,
    )
    parser.add_argument(
        "--color",
        action="store_true",
        help="colorise output to distinguish prompt and user input from generations",
    )
    parser.add_argument(
        "--keep",
        type=int,
        help="Number of tokens to keep from the initial prompt: Int (default: 0, -1 = all)",
        default=0,
    )
    parser.add_argument(
        "--shift-roped-k",
        action="store_true",
        help="Use ring-buffer and thus do not re-computing after reaching ctx_size (default: False)",
    )
    parser.add_argument(
        "--token",
        type=str,
        help="Access token ID for models that require it (LLaMa2, etc..)",
    )

    args = parser.parse_args(args_in)

    if args.model.exists():
        dir_model = args.model.as_posix()
    else:
        try:
            dir_model = snapshot_download(repo_id=str(args.model), resume_download=True, token=args.token)
        # Handles Missing token ID for gated models
        except Exception as e:
            if e.response.status_code == 401:
                print(
                    "You are required to input an access token ID for {}, please add it in option --token or download model weights locally"
                    .format(args.model))
            sys.exit(f"{e}")

    parent_path = Path(__file__).parent.absolute()
    config = AutoConfig.from_pretrained(dir_model, trust_remote_code=True)

    if args.model_type == None:
        model_type = model_maps.get(config.model_type, config.model_type)
    else:
        model_type = args.model_type

    work_path = Path(model_type + "_files")
    if not work_path.exists():
        Path.mkdir(work_path)

    # 1. convert
    path = Path(parent_path, "convert.py")
    outfile = f"gguf_{model_type}_f32" if str(args.format) == "GGUF" else f"ne_{model_type}_f32.bin"
    convert_cmd = ["python", path]
    convert_cmd.extend(["--format", str(args.format)])
    convert_cmd.extend(["--outfile", Path(work_path, outfile)])
    convert_cmd.extend(["--outtype", "f32"])
    convert_cmd.append(dir_model)
    print("Convert model ...")
    subprocess.run(convert_cmd)

    # 2. quantize
    path = Path(parent_path, "quantize.py")
    quant_file = f"gguf_{model_type}_{args.weight_dtype}.gguf" if str(
        args.format) == "GGUF" else f"ne_{model_type}_{args.weight_dtype}.bin"
    quant_cmd = ["python", path]
    quant_cmd.extend(["--model_name", model_type])
    quant_cmd.extend(["--model_file", Path(work_path, outfile + ".gguf" if str(args.format) == "GGUF" else outfile)])
    quant_cmd.extend(["--out_file", Path(work_path, quant_file)])
    quant_cmd.extend(["--weight_dtype", args.weight_dtype])
    quant_cmd.extend(["--group_size", str(args.group_size)])
    quant_cmd.extend(["--scale_dtype", args.scale_dtype])
    quant_cmd.extend(["--compute_dtype", args.compute_dtype])
    if args.use_ggml:
        quant_cmd.extend(["--use_ggml"])
    quant_cmd.extend(["--build_dir", args.build_dir])
    quant_cmd.extend(["--one_click_run", "True"])
    print("Quantize model ...")
    subprocess.run(quant_cmd)

    # 3. inference
    path = Path(parent_path, "inference.py")
    infer_cmd = ["python", path]
    infer_cmd.extend(["--model_name", model_type])
    infer_cmd.extend(["-m", Path(work_path, quant_file)])
    infer_cmd.extend(["--prompt", args.prompt])
    if args.file != None:
        infer_cmd.extend(["--file", args.file])
    infer_cmd.extend(["--n_predict", str(args.n_predict)])
    infer_cmd.extend(["--threads", str(args.threads)])
    infer_cmd.extend(["--batch_size_truncate", str(args.batch_size_truncate)])
    infer_cmd.extend(["--ctx_size", str(args.ctx_size)])
    infer_cmd.extend(["--seed", str(args.seed)])
    infer_cmd.extend(["--repeat_penalty", str(args.repeat_penalty)])
    infer_cmd.extend(["--keep", str(args.keep)])
    infer_cmd.extend(["--build_dir", args.build_dir])
    infer_cmd.extend(["--one_click_run", "True"])
    if args.shift_roped_k:
        infer_cmd.extend(["--shift-roped-k"])
    if (model_type == "baichuan" or model_type == "qwen" or model_type == "chatglm3"):
        infer_cmd.extend(["--tokenizer", dir_model])
    print("Inference model ...")
    subprocess.run(infer_cmd)


if __name__ == "__main__":
    main()
