code: def build(
      ckpt_dir: str,
      tokenizer_path: str,
      max_seq_len: int,
      max_batch_size: int,
      model_parallel_size: Optional[int] = None,
  ) -> "Llama":
      if not torch.distributed.is_initialized():
          if device == "cuda":
              torch.distributed.init_process_group("nccl")
          else:
              torch.distributed.init_process_group("gloo")
      if not model_parallel_is_initialized():
          if model_parallel_size is None:
              model_parallel_size = int(os.environ.get("WORLD_SIZE", 1))
          initialize_model_parallel(model_parallel_size)
  
      local_rank = int(os.environ.get("LOCAL_RANK", 0))
      if device == "cuda":
          torch.cuda.set_device(local_rank)
  
      # seed must be the same in all processes
      torch.manual_seed(1)
  
      if local_rank > 0:
          sys.stdout = open(os.devnull, "w")
  
      start_time = time.time()
      checkpoints = sorted(Path(ckpt_dir).glob("*.pth"))
      assert (
          len(checkpoints) > 0
      ), f"no checkpoint files found in {ckpt_dir}"
      assert model_parallel_size == len(
          checkpoints
      ), f"Loading a checkpoint for MP={len(checkpoints)} but world size is {model_parallel_size}"
      ckpt_path = checkpoints[get_model_parallel_rank()]
      checkpoint = torch.load(ckpt_path, map_location="cpu")
      with open(Path(ckpt_dir) / "params.json", "r") as f:
          params = json.loads(f.read())
  
      model_args: ModelArgs = ModelArgs(
          max_seq_len=max_seq_len,
          max_batch_size=max_batch_size,
          **params,
      )
      tokenizer = Tokenizer(model_path=tokenizer_path)
      model_args.vocab_size = tokenizer.n_words
      # support for mac
      if device == "cuda":
          if torch.cuda.is_bf16_supported():
              torch.set_default_tensor_type(torch.cuda.BFloat16Tensor)
          else:
              torch.set_default_tensor_type(torch.cuda.HalfTensor)
      else:
          torch.set_default_tensor_type(torch.HalfTensor)
      model = Transformer(model_args)
      model.load_state_dict(checkpoint, strict=False)
      model.to(device)
      print(f"Loaded in {time.time() - start_time:.2f} seconds")
  
      return Llama(model, tokenizer)
  
imports: import json
  import os
  import sys
  import time
  from pathlib import Path
  from typing import List, Literal, Optional, Tuple, TypedDict
  import torch
  import torch.nn.functional
  from fairscale.nn.model_parallel.initialize import get_model_parallel_rank, initialize_model_parallel, model_parallel_is_initialized
  from llama.model import ModelArgs, Transformer
  from llama.tokenizer import Tokenizer
label_desc: Use image node_id 1 for calling the Llama obj's underlying member method code with example usage: obj.build(ckpt_dir, tokenizer_path, max_seq_len, max_batch_size, model_parallel_size) and returns: Llama