create_config_file.py

from transformers import AutoModelWithLMHead,AutoTokenizer
import torch
from transformers import LlamaForCausalLM, AutoTokenizer, GPT2Tokenizer, GPTNeoForCausalLM, GPTNeoForCausalLM, AutoTokenizer, AutoModelWithLMHead, GPT2Tokenizer, GPTJForCausalLM
import json, torch
head_contributions = {}
model_names = ["distilgpt2", "EleutherAI/gpt-neo-125M", "gpt2", "EleutherAI/gpt-neo-1.3B", "EleutherAI/gpt-neo-2.7B"]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

models_config = {}

for model_name in model_names:

    models_config[model_name] = {}

    if model_name in ["gpt2", "gpt2-medium", "gpt2-large", "distilgpt2",  "gpt2-xl"]:
        model = AutoModelWithLMHead.from_pretrained("./saved_models/cached_models/" + model_name).to(device)
        tokenizer = AutoTokenizer.from_pretrained("./saved_models/cached_tokenizers/" + model_name, padding_side="left")   # Initialize tokenizer
        # number of heads per layer, and number of layers
        num_heads, num_layers = model.config.n_head, model.config.n_layer
        head_dim, max_length = int(model.config.n_embd/num_heads), model.config.n_positions

    elif model_name in ["EleutherAI/gpt-neo-125M", "EleutherAI/gpt-neo-1.3B", "EleutherAI/gpt-neo-2.7B"]:
        model = GPTNeoForCausalLM.from_pretrained("./saved_models/cached_models/" + model_name).to(device)
        tokenizer = GPT2Tokenizer.from_pretrained("./saved_models/cached_tokenizers/" + model_name, padding_side="left")
        num_heads, num_layers = model.config.num_heads, model.config.num_layers
        head_dim, max_length = int(model.config.hidden_size/num_heads), model.config.max_position_embeddings

    elif model_name in ["EleutherAI/gpt-j-6B"]:
        model = GPTJForCausalLM.from_pretrained("./saved_models/cached_models/" + model_name, revision="float16",torch_dtype=torch.float16,).to(device)
        tokenizer = AutoTokenizer.from_pretrained("./saved_models/cached_tokenizers/" + model_name, padding_side="left")
        num_heads, num_layers = model.config.n_head, model.config.n_layer
        head_dim, max_length = int(model.config.n_embd/num_heads), model.config.n_positions

    elif model_name in ["meta-llama/Llama-2-7b"]:
        model = LlamaForCausalLM.from_pretrained("./saved_models/cached_models/" + model_name, revision="float16",torch_dtype=torch.float16,).to(device)
        tokenizer = AutoTokenizer.from_pretrained("./saved_models/cached_tokenizers/" + model_name, padding_side="left")
        # number of heads per layer, and number of layers
        num_heads, num_layers = model.params.n_heads, model.params.n_layers
        head_dim, max_length = int(model.params.dim/num_heads), 1024

    #Assume we need 2 sections in the config file, let's call them USERINFO and SERVERCONFIG
    models_config[model_name]["num_heads"] = num_heads
    models_config[model_name]["num_layers"] = num_layers
    models_config[model_name]["head_dim"] = head_dim
    models_config[model_name]["max_length"] = max_length

json.dump(models_config, open("./model/models_config.json", "w"))