code: def restore_mistral_weight(ckpt: dict, n_head: int, n_head_kv: int):
      # see https://github.com/ggerganov/llama.cpp/blob
      # /3e73d31d9cc0232882ce61c64742aff3ecfec416/convert.py#L978
  
      for name, weight in ckpt.items():
          head, hd_size = weight.shape[0], weight.shape[1:]
          if name.endswith("attn_q.weight"):
              ckpt[name] = (
                  weight.reshape(
                      n_head, head // n_head // 2, 2, *hd_size
                  )
                  .swapaxes(1, 2)
                  .reshape(weight.shape)
              )
          elif name.endswith("attn_k.weight"):
              ckpt[name] = (
                  weight.reshape(
                      n_head_kv, head // n_head_kv // 2, 2, *hd_size
                  )
                  .swapaxes(1, 2)
                  .reshape(weight.shape)
              )
      return ckpt
  
file_path: BigDL/python/llm/src/bigdl/llm/transformers/gguf/models/mistral.py
imports: import os
  import torch
  from accelerate import init_empty_weights
  from accelerate.utils import set_module_tensor_to_device
  from tempfile import NamedTemporaryFile
  from transformers import MistralConfig, MistralForCausalLM, LlamaTokenizer
  from ..gguf import GGUFFileLoader
label_desc: Use image node_id 2 for calling a global function with example usage: restore_mistral_weight(ckpt, n_head, n_head_kv) and returns: ckpt