code: def main():
      # (0) constants
  
      models_nl = [
          "codegen-350M-nl",
          "codegen-2B-nl",
          "codegen-6B-nl",
          "codegen-16B-nl",
      ]
      models_pl = [
          "codegen-350M-multi",
          "codegen-2B-multi",
          "codegen-6B-multi",
          "codegen-16B-multi",
          "codegen-350M-mono",
          "codegen-2B-mono",
          "codegen-6B-mono",
          "codegen-16B-mono",
      ]
      models = models_nl + models_pl
  
      # (1) params
  
      parser = argparse.ArgumentParser()
      parser.add_argument(
          "--model",
          type=str,
          choices=models,
          default="codegen-350M-mono",
      )
      parser.add_argument("--device", type=str, default="cuda:0")
      parser.add_argument("--rng-seed", type=int, default=42)
      parser.add_argument(
          "--rng-deterministic", type=bool, default=True
      )
      parser.add_argument("--p", type=float, default=0.95)
      parser.add_argument("--t", type=float, default=0.2)
      parser.add_argument("--max-length", type=int, default=128)
      parser.add_argument("--batch-size", type=int, default=1)
      parser.add_argument("--no-fp16", action="store_true")
      parser.add_argument("--pad", type=int, default=50256)
      parser.add_argument(
          "--context", type=str, default="def helloworld():"
      )
      args = parser.parse_args()
  
      # (2) preamble
  
      set_env()
      set_seed(args.rng_seed, deterministic=args.rng_deterministic)
      device = torch.device(args.device)
  
      use_fp16 = True
      if args.no_fp16 or device.type == "cpu":
          use_fp16 = False
  
      if args.model.startswith("codegen-16B"):
          use_fp16 = True
  
      ckpt = f"./checkpoints/{args.model}"
  
      # (3) load
  
      with print_time("loading parameters"):
          model = create_model(ckpt=ckpt, fp16=use_fp16).to(device)
  
      with print_time("loading tokenizer"):
          if args.model in models_pl:
              tokenizer = create_custom_gpt2_tokenizer()
          else:
              tokenizer = create_tokenizer()
          tokenizer.padding_side = "left"
          tokenizer.pad_token = args.pad
  
      # (4) sample
  
      with print_time("sampling"):
          completion = sample(
              device=device,
              model=model,
              tokenizer=tokenizer,
              context=args.context,
              pad_token_id=args.pad,
              num_return_sequences=args.batch_size,
              temp=args.t,
              top_p=args.p,
              max_length_sample=args.max_length,
          )[0]
          truncation = truncate(completion)
  
          print("=" * 100)
          print(completion)
          print("=" * 100)
          print(args.context + truncation)
          print("=" * 100)
  
imports: import os
  import re
  import time
  import random
  import argparse
  import torch
  from transformers import GPT2TokenizerFast
  from jaxformer.hf.codegen.modeling_codegen import CodeGenForCausalLM
label_desc: Use image node_id 12 for calling a global function with example usage: main() without return types