code: def main(args):
      print("\nCreating retrieval dataset")
      _, _, test_dataset, code_dataset = create_dataset(
          args.data_dir, args.lang
      )
  
      test_loader, code_loader = create_loader(
          [test_dataset, code_dataset],
          [None, None],
          batch_size=[args.batch_size, args.batch_size],
          num_workers=[4, 4],
          is_trains=[False, False],
          collate_fns=[None, None],
      )
  
      tokenizer = AutoTokenizer.from_pretrained(
          args.model_name, trust_remote_code=True
      )
      model = AutoModel.from_pretrained(
          args.model_name, trust_remote_code=True
      )
      print(
          f"Loaded {args.model_name} model (#para={model.num_parameters()})"
      )
  
      print("\nStart zero-shot evaluation...")
      device = torch.device(args.device)
      model = model.to(device)
      model.eval()
  
      text_embeds = get_feats(
          model,
          tokenizer,
          test_loader,
          args.max_text_len,
          device,
          desc="Get text feats",
      )
      code_embeds = get_feats(
          model,
          tokenizer,
          code_loader,
          args.max_code_len,
          device,
          desc="Get code feats",
      )
      test_result = contrast_evaluation(
          text_embeds, code_embeds, test_loader.dataset.text2code
      )
      print(f"\n====> zero-shot test result: ", test_result)
  
      if args.local_rank in [-1, 0]:
          log_stats = {
              **{f"test_{k}": v for k, v in test_result.items()},
              "epoch": -1,
          }
  
          with open(
              os.path.join(args.output_dir, "result.txt"), "a"
          ) as f:
              f.write(json.dumps(log_stats) + "\n")
  
file_path: codet5/CodeT5+/code_retrieval/eval_contrast_retrieval.py
imports: import argparse
  import os
  import pprint
  import json
  import time
  import datetime
  import numpy
  from tqdm import tqdm
  import torch
  from transformers import AutoTokenizer, AutoModel
  from data_utils import create_dataset, create_loader
label_desc: Use image node_id 3 for calling a global function with example usage: main(args) without return types