code: def sample(
      device,
      model,
      tokenizer,
      prompt,
      pad_token_id,
      num_return_sequences=1,
      temp=0.2,
      top_p=0.95,
      max_length=2048,
      max_gen_length=128,
  ):
      input_ids = tokenizer(
          prompt,
          truncation=True,
          padding=True,
          return_tensors="pt",
      ).input_ids
  
      input_ids_len = input_ids.shape[1]
      assert input_ids_len < max_length
  
      with torch.no_grad():
          input_ids = input_ids.to(device)
          tokens = model.generate(
              input_ids,
              do_sample=True,
              num_return_sequences=num_return_sequences,
              temperature=temp,
              max_length=input_ids_len + max_gen_length,
              top_p=top_p,
              pad_token_id=pad_token_id,
              use_cache=True,
          )
          text = tokenizer.batch_decode(tokens[:, input_ids_len:, ...])
  
      return text
  
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 10 for calling a global function with example usage: sample(device, model, tokenizer, prompt, pad_token_id, num_return_sequences, temp, top_p, max_length, max_gen_length) and returns: text