code: def sample_completions(
      sample,
      device,
      model,
      tokenizer,
      n,
      t,
      p,
      pad_token_id,
      max_length,
      set_rng_seed,
      out_file,
      batch_size,
      problem_set,
      max_gen_length=256,
      prefix="# Import libraries.\n\nimport numpy as np\n\n",
  ):
      with print_time("sample completions"):
          wrap = lambda prompt: f"# {prompt}\n"
  
          for i, problem in enumerate(problem_set):
              if os.path.exists(out_file(problem["id"])):
                  print(f'skipping problem {problem["id"]}')
                  continue
  
              samples = []
  
              print("=" * 10)
              print(f'Problem {problem["id"]}')
              print("=" * 10)
  
              set_rng_seed()
  
              num_batches = n // batch_size
              remainder = n % batch_size
  
              for j in range(
                  num_batches + 1 if remainder > 0 else num_batches
              ):
                  filled_prompts = [
                      (
                          [
                              p.format(**input)
                              for p in problem["prompts"]
                          ],
                          input,
                          output,
                      )
                      for (input, output) in zip(
                          problem["inputs"], problem["outputs"]
                      )
                  ]
  
                  for k, (prompts, input, output) in enumerate(
                      filled_prompts
                  ):
                      histories = [
                          prefix
                          for _ in range(
                              batch_size
                              if j != num_batches
                              else remainder
                          )
                      ]
                      histories_full = [
                          [prefix]
                          for _ in range(
                              batch_size
                              if j != num_batches
                              else remainder
                          )
                      ]
  
                      for l, prompt in enumerate(prompts):
                          histories = [
                              h + wrap(prompt) for h in histories
                          ]
                          histories_full = [
                              h + [wrap(prompt)] for h in histories_full
                          ]
  
                          completions = sample(
                              device,
                              model,
                              tokenizer,
                              histories,
                              num_return_sequences=1,
                              top_p=p,
                              temp=t,
                              pad_token_id=pad_token_id,
                              max_length=max_length,
                              max_gen_length=problem.get(
                                  "max_gen_length", max_gen_length
                              ),
                          )
  
                          histories = [
                              h + f"{truncate(c)}\n\n"
                              for h, c in zip(histories, completions)
                          ]
                          histories_full = [
                              h + [f"{truncate(c)}\n\n"]
                              for h, c in zip(
                                  histories_full, completions
                              )
                          ]
  
                          print("-" * 10)
                          print(l)
                          print("-" * 10)
                          print(histories[0])
                          print("-" * 10)
  
                      for history, history_full in zip(
                          histories, histories_full
                      ):
                          samples.append(
                              {
                                  "id": problem["id"],
                                  "input": input,
                                  "gold_output": output,
                                  "completions": history,
                                  "prompts_completions": history_full,
                              }
                          )
  
              write_jsonl(out_file(problem["id"]), samples)
  
imports: import argparse
  import json
  import os
  from pathlib import Path
  import random
  from time import time
  import torch
  from transformers import GPT2TokenizerFast
  from jaxformer.hf.codegen.modeling_codegen import CodeGenForCausalLM
label_desc: Use image node_id 14 for calling a global function with example usage: sample_completions(sample, device, model, tokenizer, n, t, p, pad_token_id, max_length, set_rng_seed, out_file, batch_size, problem_set, max_gen_length, prefix) without return types