code: def get_model_context(
      model_key: str,
      model_name_or_path: str | None = None,
      tokenization_context: TokenizationContext | None = None,
      inference_mode: bool = True,
      use_flash_attention: bool = False,
  ) -> ModelContext:
      # `model_key` defines the model and the tokenizer to use, while `model_name_or_path`
      # defines where to load the weights. It can be from a local directory.
      assert model_key in SupportedModelKeys.all(), model_key
      if model_name_or_path is None:
          model_name_or_path = model_key
      if model_key in SupportedModelKeys.codellama_based_models():
          max_context_size = 16384
      elif model_key in SupportedModelKeys.starcoder_based_models():
          max_context_size = 8192
      else:
          assert (
              model_key
              in SupportedModelKeys.deepseekcoder_based_models()
          )
          max_context_size = 16384
      if tokenization_context is None:
          tokenization_context = TokenizationContext.from_model_key(
              model_key
          )
      # TODO: check if all these models use bfloat16
      dtype = torch.bfloat16
      other_kwargs: dict = {}
      if inference_mode:
          other_kwargs["device_map"] = "auto"
      if use_flash_attention:
          other_kwargs["use_flash_attention_2"] = True
      model = AutoModelForCausalLM.from_pretrained(
          model_name_or_path,
          torch_dtype=dtype,
          **other_kwargs,
      )
      return ModelContext(tokenization_context, model, max_context_size)
  
imports: from dataclasses import dataclass, field
  from enum import Enum
  from typing import Callable, Literal
  import torch
  from transformers import AutoModelForCausalLM, AutoTokenizer
  from transformers import GenerationConfig
  from transformers import PreTrainedModel, PreTrainedTokenizer
label_desc: Use image node_id 2 for calling a global function with example usage: get_model_context(model_key, model_name_or_path, tokenization_context, inference_mode, use_flash_attention) and returns: ModelContext