code: def __init__(
      self,
      config_cls,
      model_cls,
      model_id: str,
      revision: Optional[str] = None,
      quantize: Optional[str] = None,
      dtype: Optional[torch.dtype] = None,
      trust_remote_code: bool = False,
  ):
      global SLIDING_WINDOW
      global SLIDING_WINDOW_BLOCKS
  
      (
          self.process_group,
          rank,
          world_size,
      ) = initialize_torch_distributed()
      if torch.cuda.is_available():
          device = torch.device(f"cuda:{rank}")
          dtype = torch.float16 if dtype is None else dtype
      else:
          raise NotImplementedError(
              "FlashLlama is only available on GPU"
          )
  
      tokenizer = LlamaTokenizerFast.from_pretrained(
          model_id,
          revision=revision,
          padding_side="left",
          truncation_side="left",
          trust_remote_code=trust_remote_code,
      )
  
      config = config_cls.from_pretrained(
          model_id,
          revision=revision,
          trust_remote_code=trust_remote_code,
      )
      config.quantize = quantize
  
      # Set context windows
      if config.sliding_window is not None:
          SLIDING_WINDOW = config.sliding_window
          SLIDING_WINDOW_BLOCKS = math.ceil(
              config.sliding_window / BLOCK_SIZE
          )
  
      torch.distributed.barrier(group=self.process_group)
  
      filenames = weight_files(
          model_id, revision=revision, extension=".safetensors"
      )
      weights = Weights(
          filenames, device, dtype, process_group=self.process_group
      )
      if config.quantize in ["gptq", "awq"]:
          weights._set_gptq_params(model_id, revision)
  
      model = model_cls(config, weights)
  
      torch.distributed.barrier(group=self.process_group)
      super(BaseFlashMistral, self).__init__(
          model=model,
          tokenizer=tokenizer,
          num_layers=len(model.model.layers),
          num_kv_heads=model.model.num_key_value_heads,
          head_size=model.model.head_size,
          dtype=dtype,
          device=device,
          rank=rank,
          world_size=world_size,
          sliding_window=config.sliding_window,
      )
  
imports: import math
  import torch
  import torch.distributed
  import numpy
  from dataclasses import dataclass
  from opentelemetry import trace
  from transformers import PreTrainedTokenizerBase
  from transformers.models.llama import LlamaTokenizerFast
  from typing import Optional, Tuple, Type, List
  from text_generation_server.pb import generate_pb2
  from text_generation_server.models import FlashCausalLM
  from text_generation_server.models.flash_causal_lm import FlashCausalLMBatch, BLOCK_SIZE
  from text_generation_server.models.cache_manager import get_cache_manager
  from text_generation_server.models.custom_modeling.flash_mistral_modeling import FlashMistralForCausalLM, MistralConfig
  from text_generation_server.utils.speculate import get_speculate
  from text_generation_server.utils import initialize_torch_distributed, weight_files, Weights, HeterogeneousNextTokenChooser, StoppingCriteria
label_desc: Use image node_id 1 to create a new BaseFlashMistral object from inherited base classes: FlashCausalLM with example: obj = BaseFlashMistral(config_cls, model_cls, model_id, revision, quantize, dtype, trust_remote_code)