code: def __init__(self, config):
      self.config = config
  
      # Copy tuning parameters to C++ extension
  
      self.config.set_tuning_params()
  
      # Read tensor list from file(s)
  
      if isinstance(self.config.model_path, str):
          model_path = [self.config.model_path]
      else:
          model_path = self.config.model_path
  
      # Read tensor list from file(s), and measure layer sizes
  
      load_keys = {}
  
      decoder_size = 0
      norm_size = 0
      head_size = 0
  
      for path in model_path:
          with safe_open(path, framework="pt", device="cpu") as f:
              for key in f.keys():
                  if _skip_key(key):
                      continue
  
                  load_keys[key] = path
  
                  if key.startswith("model.layers.0."):
                      tensor_slice = f.get_slice(key)
                      shape = tensor_slice.get_shape()
                      decoder_size += math.prod(
                          shape
                      ) * _layer_dtype_size(key)
                      del tensor_slice
  
                  if key.startswith("model.norm."):
                      tensor_slice = f.get_slice(key)
                      shape = tensor_slice.get_shape()
                      norm_size += math.prod(shape) * _layer_dtype_size(
                          key
                      )
                      del tensor_slice
  
                  if key.startswith("lm_head."):
                      tensor_slice = f.get_slice(key)
                      shape = tensor_slice.get_shape()
                      head_size += math.prod(shape) * _layer_dtype_size(
                          key
                      )
                      del tensor_slice
  
      # Begin auto mapping if enabled
  
      if self.config.auto_map is not None:
          self.config.device_map.embed_tokens = "cpu"
          self.config.device_map.layers = ["cuda:0"] + ["?"] * (
              self.config.num_hidden_layers - 1
          )
  
          # Assign layers automatically
  
          device_usage = 0
          device_index = 0
          layer_index_device = 0
          max_usage = self.config.auto_map[device_index] * (1024**3)
  
          for layer in range(self.config.num_hidden_layers + 2):
              this_layer_size = decoder_size
              if layer == self.config.num_hidden_layers + 0:
                  this_layer_size = norm_size
              elif layer == self.config.num_hidden_layers + 1:
                  this_layer_size = head_size
  
              while device_usage + this_layer_size > max_usage:
                  device_index += 1
                  device_usage = 0
                  layer_index_device = 0
                  max_usage = self.config.auto_map[device_index] * (
                      1024**3
                  )
                  if device_index >= len(self.config.auto_map):
                      raise ValueError(
                          "Model too large for device allocation scheme."
                      )
  
              target = f"cuda:{device_index}"
              if layer == self.config.num_hidden_layers + 0:
                  self.config.device_map.norm = target
              elif layer == self.config.num_hidden_layers + 1:
                  self.config.device_map.lm_head = target
              else:
                  self.config.device_map.layers[
                      layer
                  ] = f"cuda:{device_index}"
  
              device_usage += this_layer_size
              layer_index_device += 1
  
      # Load up to 1 GB of tensors at a time, closing and reopening the file in between each chunk
  
      max_dq_buffer_size = 0
      tensors = {}
  
      st_mem = 0
      MAX_ST_MEM = 1024**3
      f = None
      prev_path = ""
      for key, path in load_keys.items():
          device = self.config.device_map.map(key)
  
          if f is None or st_mem > MAX_ST_MEM or path != prev_path:
              if f is not None:
                  del f
              f = safe_open(path, framework="pt", device="cpu")
              prev_path = path
              st_mem = 0
  
          tensor = f.get_tensor(key)
          size = tensor.numel() * tensor.element_size()
          st_mem += size
  
          if key.endswith(".scales"):
              tensor = tensor.half()
          if key == "lm_head.weight":
              tensor = (
                  tensor.float() if device == "cpu" else tensor.half()
              )
          if key == "model.norm.weight":
              tensor = tensor.half()
          if key.endswith(".embed_tokens.weight"):
              tensor = tensor.half()
          if key.endswith(".input_layernorm.weight"):
              tensor = tensor.half()
          if key.endswith(".post_attention_layernorm.weight"):
              tensor = tensor.half()
  
          if device == "cpu":
              keep_tensor = tensor.clone()
          else:
              keep_tensor = tensor.to(device)
          del tensor
  
          if key.endswith(".qweight"):
              max_dq_buffer_size = max(
                  max_dq_buffer_size, keep_tensor.numel() * 8
              )
  
          tensors[key] = keep_tensor
  
      del f
  
      # Head
  
      self.lm_head = nn.Linear(
          self.config.hidden_size,
          self.config.vocab_size,
          bias=False,
          device="meta",
      )
      self.lm_head.weight = nn.Parameter(tensors["lm_head.weight"])
      # self.lm_head_data = tensors["lm_head.weight"].transpose(0, 1).contiguous()
  
      # Token embeddings
  
      self.embed_tokens = nn.Embedding(
          self.config.vocab_size,
          self.config.hidden_size,
          self.config.pad_token_id,
          device="meta",
      )
      self.embed_tokens.weight = nn.Parameter(
          tensors["model.embed_tokens.weight"]
      )
      with torch.no_grad():
          self.embed_tokens.weight[self.config.pad_token_id] = 0
  
      # Norm
  
      self.norm = ExLlamaRMSNorm(
          self.config, tensors, "model.norm.weight"
      )
  
      # Prepare position embeddings for max seq length
  
      devs = self.config.device_map.get_layers_devs()
  
      self.sincos = {}
      for device in devs:
          inv_freq = 1.0 / (
              self.config.rotary_embedding_base
              ** (
                  torch.arange(
                      0, self.config.head_dim, 2, device=device
                  ).float()
                  / self.config.head_dim
              )
          )
          t = torch.arange(
              self.config.max_seq_len,
              device=device,
              dtype=torch.float32,
          )
          if self.config.compress_pos_emb != 1.0:
              t /= self.config.compress_pos_emb
  
          freqs = torch.einsum("i,j->ij", t, inv_freq)
          emb = torch.cat((freqs, freqs), dim=-1)
  
          sin = emb.sin()[None, None, :, :].half()
          cos = emb.cos()[None, None, :, :].half()
  
          self.sincos[device] = (sin, cos)
  
      # Decoder layers
  
      modules = []
      device_layer_index = [0] * len(devs)
  
      for i in range(self.config.num_hidden_layers):
          device = self.config.device_map.layers[i]
          sin, cos = self.sincos[device]
  
          layer = ExLlamaDecoderLayer(
              self.config, tensors, f"model.layers.{i}", i, sin, cos
          )
  
          modules.append(layer)
  
      self.layers = modules
  
      # Prepare CUDA buffers
  
      self.buffers = []
      for dev in self.config.device_map.get_layers_devs():
          device_buffers = {}
          self.buffers.append(device_buffers)
  
          temp_state = torch.zeros(
              (config.max_input_len, config.intermediate_size),
              dtype=torch.float16,
              device=dev,
          )
          temp_mlp = torch.zeros(
              (config.fused_mlp_thd * 2, config.intermediate_size),
              dtype=torch.float16,
              device=dev,
          )
          temp_zeros_float = torch.zeros(
              (1, 65536), dtype=torch.float32, device=dev
          )
          temp_dq = torch.zeros(
              (1, max_dq_buffer_size), dtype=torch.float16, device=dev
          )
  
          device_buffers["temp_state"] = temp_state
          device_buffers["temp_mlp"] = temp_mlp
          device_buffers["temp_zeros_float"] = temp_zeros_float
          device_buffers["temp_dq"] = temp_dq
  
          cuda_ext.exllama_ext.prepare_buffers(
              torch.device(dev),
              temp_state,
              temp_mlp,
              temp_zeros_float,
              temp_dq,
          )
  
      # Clear the cache
  
      torch.cuda.empty_cache()
  
imports: import sys
  import torch
  from torch import nn
  import torch.nn.functional
  from safetensors import safe_open
  import cuda_ext
  import json
  import math
  import gc
  from enum import Enum
label_desc: Use image node_id 1 to create a new ExLlama object  with example: obj = ExLlama(config)