docstr: Build a Llama instance by initializing and loading a pre-trained model.
  
  Args:
      ckpt_dir (str): Path to the directory containing checkpoint files.
      tokenizer_path (str): Path to the tokenizer file.
      max_seq_len (int): Maximum sequence length for input text.
      max_batch_size (int): Maximum batch size for inference.
      model_parallel_size (Optional[int], optional): Number of model parallel processes.
          If not provided, it's determined from the environment. Defaults to None.
  
  Returns:
      Llama: An instance of the Llama class with the loaded model and tokenizer.
  
  Raises:
      AssertionError: If there are no checkpoint files in the specified directory,
          or if the model parallel size does not match the number of checkpoint files.
  
  Note:
      This method initializes the distributed process group, sets the device to CUDA,
      and loads the pre-trained model and tokenizer.
docstr_tok: ["Build","a","Llama","instance","by","initializing","and","loading","a","pre-trained","model",".","Args",":","ckpt_dir","(","str",")",":","Path","to","the","directory","containing","checkpoint","files",".","tokenizer_path","(","str",")",":","Path","to","the","tokenizer","file",".","max_seq_len","(","int",")",":","Maximum","sequence","length","for","input","text",".","max_batch_size","(","int",")",":","Maximum","batch","size","for","inference",".","model_parallel_size","(","Optional","[","int","]",",","optional",")",":","Number","of","model","parallel","processes",".","If","not","provided",",","it","'s","determined","from","the","environment",".","Defaults","to","None",".","Returns",":","Llama",":","An","instance","of","the","Llama","class","with","the","loaded","model","and","tokenizer",".","Raises",":","AssertionError",":","If","there","are","no","checkpoint","files","in","the","specified","directory",",","or","if","the","model","parallel","size","does","not","match","the","number","of","checkpoint","files",".","Note",":","This","method","initializes","the","distributed","process","group",",","sets","the","device","to","CUDA",",","and","loads","the","pre-trained","model","and","tokenizer","."]
code: def build(
      ckpt_dir: str,
      tokenizer_path: str,
      max_seq_len: int,
      max_batch_size: int,
      num_gpus: int,
      model_parallel_size: Optional[int] = None,
      seed: int = 1,
  ) -> "Mixtral":
      """
      Build a Llama instance by initializing and loading a pre-trained model.
  
      Args:
          ckpt_dir (str): Path to the directory containing checkpoint files.
          tokenizer_path (str): Path to the tokenizer file.
          max_seq_len (int): Maximum sequence length for input text.
          max_batch_size (int): Maximum batch size for inference.
          model_parallel_size (Optional[int], optional): Number of model parallel processes.
              If not provided, it's determined from the environment. Defaults to None.
  
      Returns:
          Llama: An instance of the Llama class with the loaded model and tokenizer.
  
      Raises:
          AssertionError: If there are no checkpoint files in the specified directory,
              or if the model parallel size does not match the number of checkpoint files.
  
      Note:
          This method initializes the distributed process group, sets the device to CUDA,
          and loads the pre-trained model and tokenizer.
  
      """
      model_parallel_size = 1
  
      # seed must be the same in all processes
      torch.manual_seed(seed)
  
      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[0]
      with open(Path(ckpt_dir) / "params.json", "r") as f:
          params = json.loads(f.read())
  
      model_args: MixtralModelArgs = MixtralModelArgs(
          max_seq_len=max_seq_len,
          max_batch_size=max_batch_size,
          num_gpus=num_gpus,
          **params,
      )
      tokenizer = Tokenizer(model_path=tokenizer_path)
      model_args.vocab_size = tokenizer.n_words
      torch.set_default_tensor_type(torch.cuda.HalfTensor)
      model = MoETorchTransformer(model_args)
      print(
          f"=== created Mixtral 8x7B. Experts spread over {num_gpus} GPUs ==="
      )
      model_param_keys = []
      for key, value in model.named_parameters():
          model_param_keys.append(key)
  
      checkpoint = torch.load(ckpt_path, map_location="cpu")
      print(
          "Total number of model parameters:{}".format(
              len(model_param_keys)
          )
      )
      print(
          "Total number of checkpoint parameters:{}".format(
              len(checkpoint)
          )
      )
  
      model.load_state_dict(checkpoint, strict=False)
      print(f"Loaded in {time.time() - start_time:.2f} seconds")
  
      return Mixtral(model, tokenizer)
  
file_path: mixtral-kit/mixtralkit/mixtral/generation.py
imports: import json
  import os
  import sys
  import time
  from pathlib import Path
  from typing import List, Optional, Tuple, TypedDict
  import torch
  import torch.nn.functional
  from mixtralkit.layers import Tokenizer, MoETorchTransformer, MixtralModelArgs
  from mixtralkit.utils import sample_top_p
  from mixtralkit.utils.generation import CompletionPrediction, Dialog, ChatPrediction
label_desc: Use image node_id 1 for calling the Mixtral obj's underlying member method code with example usage: obj.build(ckpt_dir, tokenizer_path, max_seq_len, max_batch_size, num_gpus, model_parallel_size, seed) and returns: Mixtral