content: import json
  import os
  from abc import ABC, abstractmethod
  from pathlib import Path
  from typing import Dict, List, Optional, Tuple, Union
  
  import torch
  import torch.distributed
  from huggingface_hub import hf_hub_download
  from huggingface_hub.utils import EntryNotFoundError, LocalEntryNotFoundError
  from loguru import logger
  from safetensors import safe_open
  
  from lorax_server.utils.sources import PBASE, S3, map_pbase_model_id_to_s3
  from lorax_server.utils.torch_utils import is_fp8
  
  
  class AbstractWeights(ABC):
      @abstractmethod
      def get_slice(self, tensor_name: str) -> torch.Tensor:
          pass
  
      @abstractmethod
      def get_tensor(self, tensor_name: str) -> torch.Tensor:
          pass
  
      @abstractmethod
      def get_slice_shape(self, slice) -> torch.Size:
          pass
  
      @abstractmethod
      def has_tensor(self, tensor_name: str) -> bool:
          pass
  
      @property
      @abstractmethod
      def process_group(self):
          pass
  
      def get_shape(self, tensor_name: str) -> torch.Size:
          return self.get_slice_shape(self.get_slice(tensor_name))
  
      def get_partial_sharded(self, tensor_name: str, dim: int, range: Optional[Tuple[int, int]] = None):
          """Loads tensor with the given name and shards it along the given dimension.
  
          The optional range argument can be used to load and split on only a subset of the tensor.
          This is useful in cases where the tensor is stored as one contiguous block, but is logically
          split into different components that need to be sharded separately. For example, when storing
          QKV weights together as a single tensor on disk.
  
          Args:
              tensor_name (str): Name of the tensor to load.
              dim (int): Dimension to shard along.
              range (Optional[Tuple[int, int]]): Range of indices to load and shard as (offset, size).
          """
          world_size = self.process_group.size()
          rank = self.process_group.rank()
  
          slice_ = self.get_slice(tensor_name)
          if range is not None:
              offset, size = range
          else:
              offset = 0
              size = self.get_slice_shape(slice_)[dim]
          start, stop = get_start_stop_idxs_for_rank(offset, size, rank, world_size)
  
          if dim == 0:
              tensor = slice_[start:stop]
          elif dim == 1:
              tensor = slice_[:, start:stop]
          else:
              raise NotImplementedError("Let's make that generic when needed")
          # Special case for gptq which shouldn't convert
          # u4 which are disguised as int32
          if tensor.dtype not in [torch.int32, torch.int64, torch.float8_e4m3fn, torch.float8_e5m2]:
              tensor = tensor.to(dtype=self.dtype)
          tensor = tensor.to(device=self.device)
          return tensor
  
      def get_sharded(self, tensor_name: str, dim: int, range: Optional[Tuple[int, int]] = None):
          slice_ = self.get_slice(tensor_name)
          world_size = self.process_group.size()
          size = self.get_slice_shape(slice_)[dim] if range is None else range[1]
          assert size % world_size == 0, f"The choosen size {size} is not compatible with sharding on {world_size} shards"
          return self.get_partial_sharded(tensor_name, dim, range=range)
  
      def get_sharded_prefix(self, module_name: str, prefix: Union[str, Tuple], dim: int):
          if isinstance(prefix, str):
              return self.get_sharded(f"{prefix}.{module_name}", dim=dim)
          else:
              assert isinstance(prefix, tuple)
              assert len(prefix) == 2
              return self.get_sharded(f"{prefix[0]}.{module_name}", dim=dim, range=prefix[1])
  
      def get_sharded_list(self, module_name: str, prefixes: List[Union[str, Tuple]], dim: int):
          return [self.get_sharded_prefix(module_name, p, dim=dim) for p in prefixes]
  
      def get_multi_weights_col(self, prefixes: List[Union[str, Tuple]], quantize: str, dim: int):
          if quantize in ["gptq", "awq"]:
              try:
                  qweight = torch.cat(self.get_sharded_list("qweight", prefixes, dim=1), dim=1)
              except RuntimeError:
                  raise RuntimeError(
                      "Cannot load `gptq` weight, make sure the model is already quantized, or quantize it with `lorax-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
                  )
  
              qzeros = torch.cat(self.get_sharded_list("qzeros", prefixes, dim=1), dim=1)
              scales = torch.cat(self.get_sharded_list("scales", prefixes, dim=1), dim=1)
              if quantize == "gptq":
                  # no tensor parallelism, so remove the range if provided
                  prefixes = [p[0] if isinstance(p, tuple) else p for p in prefixes]
                  w = [self.get_tensor(f"{p}.g_idx") for p in prefixes]
                  for w2 in w[1:]:
                      torch.testing.assert_close(w2, w[0])
                  g_idx = w[0]
              else:
                  g_idx = None
  
              bits, groupsize = self._get_bits_and_groupsize()
              weight = (qweight, qzeros, scales, g_idx, bits, groupsize, False)
          else:
              weight_list = self.get_sharded_list("weight", prefixes, dim=0)
              if is_fp8(quantize) and weight_list[0].dtype == torch.float8_e4m3fn:
                  # Since there is no kernel for concatenating two tensors in PyTorch
                  # for fp8 datatypes, we have to cast to fp16, concat, cast back to fp8
                  fp16_weight_list = [w.to(torch.float16) for w in weight_list]
                  weight = torch.cat(fp16_weight_list, dim=dim).to(torch.float8_e4m3fn)
                  input_scale = None
                  if self.has_tensor(f"{prefixes[0]}.input_scale"):
                      # if the layers are being fused, then they have the same inputs
                      # hence their input scales will have to be the same so we pick the first one
                      input_scale = self.get_tensor(f"{prefixes[0]}.input_scale", use_self_dtype=False)
                  weight_scale_list = [self.get_tensor(f"{p}.weight_scale", use_self_dtype=False) for p in prefixes]
                  if len(weight_scale_list[0].shape) > 1:
                      weight_scale_list = self.get_sharded_list("weight_scale", prefixes, dim=0)
                  else:
                      weight_scale_list = [si.repeat(wi.shape[dim]) for si, wi in zip(weight_scale_list, weight_list)]
                  # weight scales are in fp32 already so no problem with concatenating them
                  weight_scale = torch.cat(weight_scale_list, dim=0)
                  return weight, input_scale, weight_scale
              weight = torch.cat(weight_list, dim=dim)
  
          return weight
  
      def get_multi_weights_row(self, prefix: str, quantize: str):
          if quantize == "gptq":
              use_exllama = True
              bits, groupsize = self._get_bits_and_groupsize()
  
              if bits != 4:
                  use_exllama = False
  
              if self.process_group.size() > 1:
                  g_idx = self.get_tensor(f"{prefix}.g_idx")
                  if g_idx is not None:
                      if (
                          not torch.equal(
                              g_idx.cpu(),
                              torch.tensor(
                                  [i // groupsize for i in range(g_idx.shape[0])],
                                  dtype=torch.int32,
                              ),
                          )
                          and not (g_idx == 0).all()
                      ):
                          # Exllama implementation does not support row tensor parallelism with act-order, as
                          # it would require to reorder input activations that are split unto several GPUs
                          use_exllama = False
  
              try:
                  qweight = self.get_sharded(f"{prefix}.qweight", dim=0)
              except RuntimeError:
                  raise RuntimeError(
                      "Cannot load `gptq` weight, make sure the model is already quantized, or quantize it with `lorax-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
                  )
  
              from lorax_server.utils.layers import HAS_EXLLAMA
  
              if use_exllama:
                  if not HAS_EXLLAMA:
                      logger.warning(
                          "Exllama GPTQ cuda kernels (which are faster) could have been used, but are not currently installed, try using BUILD_EXTENSIONS=True"
                      )
                      use_exllama = False
                  else:
                      logger.info("Using exllama kernels")
  
              if use_exllama:
                  if groupsize >= 0:
                      # Exllama reorders the weights in advance and the activations on the fly, thus
                      # the scales and zero-points do not need to be reordered.
                      qzeros = self.get_sharded(f"{prefix}.qzeros", dim=0)
                      scales = self.get_sharded(f"{prefix}.scales", dim=0)
                  else:
                      qzeros = self.get_tensor(f"{prefix}.qzeros")
                      scales = self.get_tensor(f"{prefix}.scales")
  
                  # For tp > 1, at this point we know we do not use act-order
                  if self.process_group.size() == 1:
                      g_idx = self.get_tensor(f"{prefix}.g_idx")
                  else:
                      g_idx = None
              else:
                  # The triton kernel reorders the scales/zero points instead of the weight/activation.
                  # Thus, each rank needs the full qzeros/scales.
                  qzeros = self.get_tensor(f"{prefix}.qzeros")
                  scales = self.get_tensor(f"{prefix}.scales")
                  g_idx = self.get_sharded(f"{prefix}.g_idx", dim=0)
  
              weight = (qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama)
          elif quantize == "awq":
              bits, groupsize = self._get_bits_and_groupsize()
  
              try:
                  qweight = self.get_sharded(f"{prefix}.qweight", dim=0)
              except RuntimeError:
                  raise RuntimeError("Cannot load `awq` weight, make sure the model is already quantized")
  
              qzeros = self.get_sharded(f"{prefix}.qzeros", dim=0)
              scales = self.get_sharded(f"{prefix}.scales", dim=0)
              g_idx = None
              use_exllama = False
              weight = (qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama)
          else:
              weight = self.get_sharded(f"{prefix}.weight", dim=1)
              if is_fp8(quantize) and weight.dtype == torch.float8_e4m3fn:
                  # weight_scale could be a tensor but if we're sharding row-wise then no
                  # need to shard the weight_scale as its row dimension would be 1
                  weight_scale = self.get_tensor(f"{prefix}.weight_scale", use_self_dtype=False)
                  input_scale = None
                  if self.has_tensor(f"{prefix}.input_scale"):
                      input_scale = self.get_tensor(f"{prefix}.input_scale", use_self_dtype=False)
                  return weight, input_scale, weight_scale
          return weight
  
      def _get_bits_and_groupsize(self) -> Tuple[int, int]:
          try:
              bits = self.config.quantization_config["bits"]
              groupsize = self.config.quantization_config["group_size"]
          except KeyError:
              # be compatible with old hehavior for gptq
              try:
                  bits = self.config.quantization_config["gptq_bits"]
                  groupsize = self.config.quantization_config["gptq_groupsize"]
              except KeyError:
                  try:
                      bits = self.get_tensor("gptq_bits").item()
                      groupsize = self.get_tensor("gptq_groupsize").item()
                  except Exception as e:
                      raise e
  
          return bits, groupsize
  
  
  class InMemoryWeights(AbstractWeights):
      def __init__(
          self,
          weights: Dict[str, torch.Tensor],
          device: torch.device,
          dtype: torch.dtype,
          process_group: torch.distributed.ProcessGroup,
      ):
          self.weights = weights
          self.device = device
          self.dtype = dtype
          self._process_group = process_group
  
      def get_slice(self, tensor_name: str) -> torch.Tensor:
          return self.get_tensor(tensor_name)
  
      def get_tensor(self, tensor_name: str) -> torch.Tensor:
          tensor = self.weights[tensor_name]
          return load_module_weight(tensor_name, tensor, self.device, self.dtype)
  
      def get_slice_shape(self, slice) -> torch.Size:
          return slice.shape
  
      def has_tensor(self, tensor_name: str) -> bool:
          return tensor_name in self.weights
  
      @property
      def process_group(self):
          return self._process_group
  
  
  class Weights(AbstractWeights):
      """
      A class representing weights for a model.
  
      Args:
          filenames (List[Path]): List of file paths containing the weights.
          device: The device to load the weights onto.
          dtype: The data type to convert the weights to.
          process_group: The process group for distributed training.
          aliases (Optional[Dict[str, List[str]]]): Dictionary of aliases for weight names.
          merged_weight_filenames (Optional[List]): List of file paths containing merged weights.
  
      Attributes:
          aliases (Dict[str, List[str]]): Dictionary of aliases for weight names.
          routing (Dict[str, str]): Dictionary mapping weight names to file paths.
          device: The device to load the weights onto.
          dtype: The data type of the weights.
          process_group: The process group for distributed training.
          _handles (Dict[str, Any]): Dictionary of file handles for opened weight files.
      """
  
      def __init__(
          self,
          filenames: List[Path],
          device,
          dtype,
          process_group,
          aliases: Optional[Dict[str, List[str]]] = None,
          merged_weight_filenames: Optional[List] = None,
      ):
          # routes to adapter files take precedence over routes to main model files
          # to ensure that adapter weights are loaded instead of main model weights
          routing = {}
          if merged_weight_filenames is not None:
              for filename in merged_weight_filenames:
                  with safe_open(filename, framework="pytorch") as f:
                      for k in f.keys():
                          if k in routing:
                              raise RuntimeError(
                                  f"Key {k} was found in multiple adapter files: {filename} and {routing[k]}"
                              )
                          routing[k] = filename
  
          # set of keys that point to adapter files. Duplicates for these keys found
          # in main model files will be overridden.
          adapter_routes = set(routing.keys())
  
          for filename in filenames:
              with safe_open(filename, framework="pytorch") as f:
                  for k in f.keys():
                      if k in adapter_routes:
                          logger.debug(f"Overriding main model weights with adapter weights for key: {k}")
                      elif k in routing:
                          raise RuntimeError(
                              f"Key {k} was found in multiple non-adapter files: {filename} and {routing[k]}"
                          )
                      else:
                          routing[k] = filename
  
          if aliases is None:
              aliases = {}
          self.aliases = aliases
          self.routing = routing
          self.device = device
          self.dtype = dtype
          self._process_group = process_group
          self._handles = {}
  
      @property
      def process_group(self):
          return self._process_group
  
      def _get_handle(self, filename):
          if filename not in self._handles:
              f = safe_open(filename, framework="pytorch")
              self._handles[filename] = f
  
          return self._handles[filename]
  
      def get_filename(self, tensor_name: str) -> (str, str):
          filename = self.routing.get(tensor_name, None)
          if filename is None:
              aliases = self.aliases.get(tensor_name, [])
              for alias in aliases:
                  filename = self.routing.get(alias, None)
                  if filename is not None:
                      return str(filename), alias
              raise RuntimeError(f"weight {tensor_name} does not exist")
          return str(filename), tensor_name
  
      def get_slice(self, tensor_name: str):
          filename, tensor_name = self.get_filename(tensor_name)
          f = self._get_handle(filename)
          slice_ = f.get_slice(tensor_name)
          return slice_
  
      def get_slice_shape(self, slice) -> torch.Size:
          return slice.get_shape()
  
      def get_tensor(self, tensor_name: str, use_self_dtype: bool = True):
          filename, tensor_name = self.get_filename(tensor_name)
          f = self._get_handle(filename)
          tensor = f.get_tensor(tensor_name)
          # Special case for gptq which shouldn't convert
          # u4 which are disguised as int32
          if tensor.dtype not in [torch.int32, torch.int64, torch.float8_e4m3fn, torch.float8_e5m2]:
              if use_self_dtype:
                  tensor = tensor.to(dtype=self.dtype)
          tensor = tensor.to(device=self.device)
          return tensor
  
      def has_tensor(self, tensor_name: str) -> bool:
          filename = self.routing.get(tensor_name, None)
          if filename is None:
              aliases = self.aliases.get(tensor_name, [])
              for alias in aliases:
                  filename = self.routing.get(alias, None)
                  if filename is not None:
                      return True
              return False
          return True
  
      def _set_config(self, model_id, config):
          self.config = config
  
          if not hasattr(self.config, "quantization_config"):
              # fill from other config file
              filename = "quantize_config.json"
              try:
                  if os.path.exists(os.path.join(model_id, filename)):
                      filename = os.path.join(model_id, filename)
                  else:
                      filename = hf_hub_download(model_id, filename=filename)
                  with open(filename, "r") as f:
                      data = json.load(f)
                  self.config.quantization_config = data["quantization_config"]
              except Exception:
                  filename = "quant_config.json"
                  try:
                      if os.path.exists(os.path.join(model_id, filename)):
                          filename = os.path.join(model_id, filename)
                      else:
                          filename = hf_hub_download(model_id, filename=filename)
                      with open(filename, "r") as f:
                          data = json.load(f)
                      self.config.quantization_config = data["quantization_config"]
                  except Exception:
                      pass
  
  
  def get_start_stop_idxs_for_rank(offset, size, rank, world_size):
      block_size = size // world_size
      start = offset + rank * block_size
      stop = offset + (rank + 1) * block_size
      return start, stop
  
  
  def shard_on_dim(t: torch.Tensor, dim: int, process_group: torch.distributed.ProcessGroup):
      world_size = process_group.size()
      rank = process_group.rank()
  
      size = t.shape[dim]
      start, stop = get_start_stop_idxs_for_rank(0, size, rank, world_size)
  
      if dim == 0:
          tensor = t[start:stop]
      elif dim == 1:
          tensor = t[:, start:stop]
      else:
          raise NotImplementedError("Let's make that generic when needed")
  
      return tensor
  
  
  def download_weights(
      model_id: str,
      revision: Optional[str] = None,
      extension: str = ".safetensors",
      auto_convert: bool = True,
      source: str = "hub",
      api_token: Optional[str] = None,
      embedding_dim: Optional[int] = None,
  ):
      # Import here after the logger is added to log potential import exceptions
      from lorax_server import utils
      from lorax_server.utils import sources
  
      if source == PBASE:
          # TODO(travis): move this into `model_source` to handle behind the abstraction
          api_token = api_token or os.environ.get("PREDIBASE_API_TOKEN")
          model_id = map_pbase_model_id_to_s3(model_id, api_token)
          source = S3
  
      model_source = sources.get_model_source(source, model_id, revision, extension, api_token, embedding_dim)
  
      # Test if files were already download
      try:
          model_source.weight_files()
          logger.info("Files are already present on the host. " "Skipping download.")
          return
      # Local files not found
      except (LocalEntryNotFoundError, FileNotFoundError):
          pass
  
      is_local_model = (Path(model_id).exists() and Path(model_id).is_dir()) or os.getenv(
          "WEIGHTS_CACHE_OVERRIDE", None
      ) is not None
  
      if not is_local_model:
          # TODO: Combine into class that takes the source as input
          # Try to download weights from the hub
          try:
              model_source.download_model_assets()
              return
          # No weights found on the hub with this extension
          except EntryNotFoundError as e:
              # Check if we want to automatically convert to safetensors or if we can use .bin weights instead
              if not extension == ".safetensors" or not auto_convert:
                  raise e
  
      # Try to see if there are local pytorch weights
      try:
          # Get weights for a local model, a hub cached model and inside the WEIGHTS_CACHE_OVERRIDE
          local_pt_files = model_source.weight_files(extension=".bin")
  
      # No local pytorch weights
      except LocalEntryNotFoundError:
          if extension == ".safetensors":
              logger.warning(
                  f"No safetensors weights found for model {model_id} at revision {revision}. "
                  f"Downloading PyTorch weights."
              )
  
          # Try to see if there are pytorch weights on the hub
          pt_filenames = model_source.remote_weight_files(extension=".bin")
          # Download pytorch weights
          local_pt_files = model_source.download_weights(pt_filenames)
  
      if auto_convert:
          logger.warning(
              f"No safetensors weights found for model {model_id} at revision {revision}. "
              f"Converting PyTorch weights to safetensors."
          )
  
          # Safetensors final filenames
          local_st_files = [p.parent / f"{p.stem.lstrip('pytorch_')}.safetensors" for p in local_pt_files]
          try:
              import transformers
              from transformers import AutoConfig
  
              config_path = sources.get_config_path(model_id, source)
              config = AutoConfig.from_pretrained(
                  config_path,
                  revision=revision,
              )
              architecture = config.architectures[0]
  
              class_ = getattr(transformers, architecture)
  
              # Name for this varible depends on transformers version.
              discard_names = getattr(class_, "_tied_weights_keys", [])
              discard_names.extend(getattr(class_, "_keys_to_ignore_on_load_missing", []))
  
          except Exception:
              discard_names = []
          # Convert pytorch weights to safetensors
          utils.convert_files(local_pt_files, local_st_files, discard_names)
  
  
  def load_module_weight(name: str, module: Union[torch.Tensor, str], device, dtype):
      if isinstance(module, torch.Tensor):
          return module.to(device, dtype)
  
      if isinstance(device, torch.device):
          if device.type == "cuda":
              device = device.index
          elif device.type == "cpu":
              device = "cpu"
  
      # module would be just the filename if lazy loading happened before
      with safe_open(module, framework="pt", device=device) as f:
          return f.get_tensor(name).to(dtype)
  