code: def main():
      parser = argparse.ArgumentParser()
  
      parser.add_argument("--input_path", type=str, default="")
      parser.add_argument("--output_path", type=str, default="")
  
      parser.add_argument("--model_name", type=str, default="")
      parser.add_argument("--model_arch", type=str, default="opt")
      parser.add_argument(
          "--cache_dir", type=str, default="../../checkpoint/"
      )
  
      parser.add_argument("--heavy_ratio", type=float, default=0.1)
      parser.add_argument("--recent_ratio", type=float, default=0.1)
      parser.add_argument("--enable_small_cache", action="store_true")
  
      parser.add_argument("--sample_num", type=int, default=1000)
  
      parser.add_argument("--k", type=int, default=0)
      parser.add_argument(
          "--seed",
          type=int,
          default=42,
          help="random seed for initialization",
      )
      parser.add_argument(
          "--no_cuda",
          action="store_true",
          help="Avoid using CUDA when available",
      )
      parser.add_argument(
          "--fp16",
          action="store_true",
          help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
      )
      args = parser.parse_args()
  
      args.device = torch.device(
          "cuda"
          if torch.cuda.is_available() and not args.no_cuda
          else "cpu"
      )
      args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
  
      logger.warning(
          f"device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fp16}"
      )
      set_seed(args)
  
      model_name = args.model_name
      input_path = args.input_path
      output_path = args.output_path
  
      config = AutoConfig.from_pretrained(
          model_name, cache_dir=args.cache_dir
      )
      tokenizer = AutoTokenizer.from_pretrained(
          model_name, use_fast=True, cache_dir=args.cache_dir
      )
      model = AutoModelForCausalLM.from_pretrained(
          model_name, cache_dir=args.cache_dir
      )
  
      if args.enable_small_cache:
          print("Enable Small Cache Size")
          config.heavy_ratio = args.heavy_ratio
          config.recent_ratio = args.recent_ratio
          checkpoint = copy.deepcopy(model.state_dict())
          model = ENABLE_Heavy_Hitter_FUNCTIONS[args.model_arch](
              model, config
          )
          model.load_state_dict(checkpoint)
  
      model.half().eval().cuda()
      logger.info(args)
  
      requests = []
      with open(input_path, "r") as f:
          for line in f:
              if line.strip() != "":
                  requests.append(json.loads(line))
  
      print(len(requests))
      if args.sample_num < len(requests):
          print("Sample {} Examples".format(args.sample_num))
      requests = requests[: args.sample_num]
  
      results = []
      with torch.no_grad():
          for request in tqdm.tqdm(requests):
              request = request["request"]
              result = {"request": request, "result": {}}
              prompt = request["prompt"]
              temperature = request["temperature"]
              stop = request["stop"]
  
              input_ids = tokenizer(
                  prompt, add_special_tokens=False, return_tensors="pt"
              ).input_ids.to(model.device)
  
              output_sequences = model.generate(
                  input_ids=input_ids,
                  max_length=request["max_tokens"] + len(input_ids[0]),
                  temperature=temperature,
                  top_k=args.k,
                  top_p=request["top_p"],
                  do_sample=True,
                  num_return_sequences=request["n"],
                  return_dict_in_generate=True,
                  output_scores=True,
              )
  
              for name, m in model.named_modules():
                  if isinstance(m, TAGET_MODULE[args.model_arch]):
                      m._reset_masks()
  
              tokens = tokenizer.convert_ids_to_tokens(
                  output_sequences["sequences"].squeeze(0)
              )[len(input_ids[0]) :]
              logprobs = [
                  logits.log_softmax(dim=-1).max().item()
                  for logits in output_sequences["scores"]
              ]
              top_logprobs = [{i: v for i, v in zip(tokens, logprobs)}]
  
              generate_text = tokenizer.decode(
                  output_sequences["sequences"].squeeze(0)[
                      len(input_ids[0]) :
                  ]
              )
              generate_text = generate_text[
                  : generate_text.find(stop[0])
              ]
  
              result["result"] = {
                  "choices": [
                      {
                          "text": generate_text,
                          "logprobs": {
                              "tokens": tokens,
                              "token_logprobs": logprobs,
                              "top_logprobs": top_logprobs,
                              "text_offset": [],
                          },
                          "finish_reason": "length",
                      }
                  ],
                  "request_time": {"batch_time": 0, "batch_size": 1},
              }
  
              results.append(result)
  
      with open(output_path, "w") as f:
          for result in results:
              f.write(json.dumps(result) + "\n")
  
file_path: H2O/h2o_hf/run_helm.py
imports: import argparse
  import logging
  import numpy
  import torch
  import json
  import tqdm
  import copy
  from transformers import CTRLLMHeadModel, CTRLTokenizer, GPT2LMHeadModel, GPT2Tokenizer, OpenAIGPTLMHeadModel, OpenAIGPTTokenizer, TransfoXLLMHeadModel, TransfoXLTokenizer, XLMTokenizer, XLMWithLMHeadModel, XLNetLMHeadModel, XLNetTokenizer
  from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
  from utils_hh.modify_llama import convert_kvcache_llama_heavy_recent, LlamaAttention_heavy_hitter
  from utils_hh.modify_gptneox import convert_kvcache_gpt_neox_heavy_recent, GPTNeoXAttention_Mask
  from utils_hh.modify_opt import convert_kvcache_opt_heavy_recent, OPTAttention_Mask
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types