import json
import os

from rtp_llm.tools.fake_model_base import *


def save_config_func(
    model_type: str,
    dest_path: str,
    layer_num: int,
    head: int,
    head_kv_num: int,
    head_size: int,
    ffn_hidden_size: int,
    ffn_inter_padding_size: int,
    vocab_size: int,
):
    config = {
        "architectures": ["GPTNeoXForCausalLM"],
        "auto_map": {"AutoModel": "model.GPTNeoXForCausalLM"},
        "bos_token_id": 2,
        "eos_token_id": 2,
        "hidden_act": "gelu",
        "hidden_size": head * head_size,
        "initializer_range": 0.02,
        "intermediate_size": ffn_hidden_size,
        "layer_norm_eps": 1e-05,
        "max_position_embeddings": 2048,
        "model_type": model_type,
        "num_attention_heads": head,
        "num_hidden_layers": layer_num,
        "rotary_emb_base": 10000,
        "rotary_pct": 1.0,
        "tie_word_embeddings": False,
        "torch_dtype": "float16",
        "transformers_version": "4.27.1",
        "use_cache": True,
        "use_parallel_residual": False,
        "vocab_size": vocab_size,
        "attention_dropout": 0,
        "hidden_dropout": 0,
        "classifier_dropout": 0.1,
        "rope_scaling": {"type": "dynamic", "factor": 2},
    }
    # save to config.json
    json.dump(config, open(os.path.join(dest_path, "config.json"), "w"), indent=2)


def fake_gpt_neox():
    default_config = DefaultModelConfig()
    default_config.layer_num = 2
    default_config.head_num = 4
    default_config.head_kv_num = 4
    default_config.head_size = 128
    default_config.ffn_hidden_size = (
        4 * default_config.head_size * default_config.head_num
    )
    default_config.ffn_inter_padding_size = default_config.ffn_hidden_size
    default_config.ffn_gate_active = False
    default_config.ffn_w1_w3_independ = False
    default_config.vocab_size = 250752

    fake_model("gpt_neox", default_config, save_config_func)


if __name__ == "__main__":
    fake_gpt_neox()
